Sorry I was not clear enough in the previous post.
I meant that if only a state dict is given and we don’t know the dimension of the output (i.e. the number of classes), then it is not possible to load the state dict.
Let me give an example (with ResNet, because that’s what I’m using).
Suppose that that the dataset has changed and now we have new_num_classes = 100 classes, but we don’t know the previous old_num_classes. All the information we have is state_dict.pth. To define a model for training, I use the following code.
# Create the model and change the dimension of the output
model = torchvision.models.resnet152()
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, new_num_classes)
# Load the pre-trained weights, only works if the dimensions are the same
model.load_state_dict(torch.load('state_dict.pth'))

If old_num_classes != new_num_classes then the last line will cause an error.
In conclusion, I have to either save the whole model using torch.save(the_model, PATH), or save only the state dict and the number of classes.