Unable to load model state after modifications (missing key(s) in state_dict)

You could filter out the unnecessary parameters in your state_dict.
Have a look at this post.