Hello!
I am currently experimenting with a network to perform multitask classification on faces.
I evaluate the model every 50 steps (switching between model.eval() and model.train() ) and obtain values like:
  "Task 1":   {"MAE": 4.9965996742248535, "RMSE": 7.064884662628174},
  "Task 2":   {"ACC": 0.9571999907493591}

After training is done, I save the model’s state dict by calling:
torch.save(model.state_dict(), os.path.join(PATH, 'model-{}.pt'.format(epoch))) 

The problem is, when loading the model in a jupyter notebook for further testing and calling the evaluation function in the exact same way as called in the training script I get the results:
"Task 1":   {'MAE': 10.82919979095459, 'RMSE': 13.656017303466797},
"Task 2":   {'ACC': 0.5248000025749207}

Which looks to me as if the loaded weights were never trained at all.
I tried training again multiple times and the results are the same. As a hail mary, I tried saving the entire model (and not just the state dict) and to my surprise it worked! The evaluation results were consistent.
Does anyone have any idea of what I’m doing wrong, and how I can get good results from the state dict so I don’t rely on the entire saved model? Thanks a lot!