A way to fix your issue might be hidden here: https://pytorch.org/tutorials/recipes/recipes/warmstarting_model_using_parameters_from_a_different_model.html 2
Maybe it’s enough to do something like:
net = MyNet()
net.load_state_dict(torch.load(PATH), strict=False)