使用 torch.load 时出现运行时错误“存储大小错误：”

如何解决使用 torch.load 时出现运行时错误“存储大小错误：”

我在调用 torch.load("pthfilename") 时收到“RuntimeError storage has wrong size”。我的模型是在多个 GPU 上训练的，我使用以下代码保存了模型：

os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4,5,6,7"
device = torch.device(arg.local_rank)
net = Net().to(device)
net = torch.nn.parallel.distributedDataParallel(net,device_ids=[arg.local_rank])
torch.save(net.state_dict(),"0.pth"))

错误是：

Traceback (most recent call last):
  File "/root/PycharmProjects/test.py",line 8,in <module>
    model_dict = torch.load("0.pth")
  File "torch/serialization.py",line 529,in load
    return _legacy_load(opened_file,map_location,pickle_module,**pickle_load_args)
  File "torch/serialization.py",line 709,in _legacy_load
    deserialized_objects[key]._set_from_file(f,offset,f_should_read_directly)
RuntimeError: storage has wrong size: expected -4916312287391674656 got 24