Error in loading the model weights on colab even though no changes to the model

Hey guys, I have trained a model on a remote workstation and pulled its weights (.pt file) onto my local machine. I am using colab to test the model however I face the following error:

RuntimeError: Error(s) in loading state_dict for CentralizedActorCritic:
size mismatch for value1.weight: copying a param with shape torch.Size([256, 6]) from checkpoint, the shape in current model is torch.Size([256, 18]).
size mismatch for policy1.weight: copying a param with shape torch.Size([256, 6]) from checkpoint, the shape in current model is torch.Size([256, 18]).

The model specifications on both colab and the remote workstation are the same. Basically, I am running the exact same code but unable to load the trained weights on colab.
Thanks :smiley:

Could you create the model locally and on Colab before loading the state_dict and print the shapes of these conflicting parameters?
Are you using some conditions to create the model, which might differ on the platforms?

Hey, I got it solved. It was some condition to do with the creation of the model! Thanks :smiley: