mobile_gallery_pytorch_02_model.py

import torch
import torch.nn as nn
from torchvision import models

# Defining the model
model = models.resnet34(pretrained = True)

# The original architecture of resnet34 has 1000 neurons(corresponding to 1000 classes on which it was originally trained on) in the final layer.
# So we need to change the final layer according to the number of classes that we have in our dataset
print('model.fc before: ', model.fc)
model_fc_in_features = model.fc.in_features
model.fc = nn.Linear(model_fc_in_features, len(dataset.classes))
print('model.fc after : ', model.fc)

mobile_gallery_pytorch_02_model_output.py

model.fc before:  Linear(in_features=512, out_features=1000, bias=True)
model.fc after :  Linear(in_features=512, out_features=6, bias=True)