Pytorch nn.parallel.DistributedDataParallel model load
Model saved with:
torch.distributed.init_process_group(backend=&quot;nccl&quot;)
local_rank = torch.distributed.get_rank()
torch.cuda.set_device(local_rank)
device = torch.device(&quot;cuda&quot;, local_rank)
save_model = f'./model'
Path(save_model).mkdir(parents=True, exist_ok=True)
net = Net(args)  # .to(device)
model_name = f&quot;{save_model}/net.pt&quot;
torch.save(net.state_dict(), model_name)   #
model = Model(net, args).to(device)
model_name = f&quot;{save_model}/model.pt&quot;

if torch.cuda.device_count() &gt; 1:
    model = nn.parallel.DistributedDataParallel(model, device_ids=[local_rank],
                                                output_device=local_rank)
model.module.fit(tr_data, val_data, args)
torch.save(maml, model_name) 

I tried to load the model with:
save_model = f'./model'
net = Net(args)  # .to(device)
model_name = f&quot;{save_model}/net.pt&quot;
net.load_state_dict(
    torch.load(model_name, map_location=torch.torch.device(&quot;cpu&quot;)))
maml = Model(net, args).to(device)
model_name = f&quot;{save_model}/model.pt&quot;
maml = torch.load(model_name, map_location=torch.torch.device(
    &quot;cuda&quot; if torch.cuda.is_available() else &quot;cpu&quot;))  # .load_state_dict

The &quot;net&quot; can be loaded successfully, but I got the error when loading &quot;model&quot;:
File &quot;D:\Research\Traffic Prediction\maml\venv\lib\site-packages\torch\serialization.py&quot;, line 607, in load return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
File &quot;D:\Research\Traffic Prediction\maml\venv\lib\site-packages\torch\serialization.py&quot;, line 882, in _load result = unpickler.load()
TypeError: () missing 1 required positional argument: 'ddp_join_throw_on_early_termination'
Any input would be really appreciated.

In general, with PyTorch’s DistributedDataParallel, same model is kept across all nodes (as it’s ‘synchronised’ during backpropagation).
Best way to save it is to just save the model instead of the whole DistributedDataParallel (usually on main node or multiple if possible node failure is a concern):
# or not only local_rank 0
if local_rank == 0:
    torch.save(model.module.cpu(), path)

Please notice, if your model is wrapped within DistributedDataParallel the model you are after is kept within module attribute.
Another thing - cast your model to CPU, no mapping will be necessary in such case (as you might use multiple GPUs and you would have to map it appropriately on other devices, which might not have GPU).
