Your Problem lies within your saving/loading code:
You save your model via torch.save(net, save_path) where net is an instance of nn.DataParallel
I would recommend to change a few things:


You should only save the state dict instead of your model. See this 142 post for further details


You should not save the stat dict of your DataParallel instance but your model’s state dict since the data parallel is simply a wrapper and also contains things like the used GPUs and copies of your model. You can access the model by calling net.module just as @JuanFMontesinos mentioned.


Taking these steps into account you would save your model with code like this:
net = nn.DataParallel(net)
    .....
    torch.save(net.module.state_dict(), save_path)

and load your model like this:
model = YourNetworkClass() # create an instance of your network
model.load_state_dict(torch.load(save_path))

If you want to infer on multiple GPUs or continue training on multiple GPUs you would have to wrap your model again with nn.DataParallel.
Also a good practice would be to move the model to cpu before saving it’s state_dict and move it back to GPU afterwards. This way the state dict will also be loaded to CPU instead of being loaded to GPU directly (which happens if you save the state_dict of a GPU-Model since you would save CUDA-Tensors). Loading a model on CPU is better practice since you could also deploy it to machines wich aren’t CUDA-capable.