Thanks, that resolved the issue.
However, I now get a KeyError when the parameters load:

KeyError: ‘unexpected key “module.feature_extractor.conv_block1.0.weight” in state_dict’

feature_extractor is a custom module that is used as part of MyModel and conv_block1 is a nn.Sequential class member of feature_extractor. Why can’t those weights be loaded? Again I saved the model as:
torch.save(model.state_dict(), save_path)