Hi @ptrblck. This worked perfectly. Apparently there were some extra keys that I didn’t take note of, so PyTorch can’t match them. However, I thought of a cleaner solution. Can’t I just freeze the classification layer when fine-tuning? It is the last layer so backprop wouldn’t be cut in the middle:
for p in model.classification.parameters():
    p.requires_grad = False