@ptrblck  Thanks a lot for this. It works after your suggestion. I will create a GitHub issue later.
I have a further question regarding saving and loading wrapped model.
In this project, I will need to save the trained model and load it for use later.
If I wrap the model like this:
wrapped_contact_net=torch.nn.DataParallel(contact_net)

I am supposed to save it like this?
        try:
            state_dict = wrapped_contact_net.module.state_dict()
        except AttributeError:
            state_dict = wrapped_contact_net.state_dict()
        torch.save(state_dict, model_path)

If I want to retrain the trained model, do I have to wrap the model name before I load?
wrapped_contact_net=torch.nn.DataParallel(contact_net)
wrapped_contact_net.load_state_dict(torch.load(model_path,map_location=device))

or wrap the model after I load?
contact_net.load_state_dict(torch.load(model_path,map_location=device))
wrapped_contact_net=torch.nn.DataParallel(contact_net)

Thanks in advance.