@ptrblck Thanks,I am OK with loading now. As for saving, do you mean I should save with the original model name, like:
state_dict = contact_net.module.state_dict()

instead of using the wrapped model:
state_dict = wrapped_contact_net.module.state_dict()