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?