如果保存模型在 GPU 上，加载模型在 CPU 上，那么我们的保存与加载的代码应该如下：
# Save on GPU device = torch.device("cuda") model.to(device) torch.save(model.state_dict(), PATH) # Load on CPU device = torch.device('cpu') model = Model(*args, **kwargs) model.load_state_dict(torch.load(PATH, map_location=device))如果保存模型在 GPU 上，加载模型在 GPU 上，那么我们的保存与加载的代码应该如下：
device = torch.device("cuda") model.to(device) torch.save(model.state_dict(), PATH) model = Model(*args, **kwargs) model.load_state_dict(torch.load(PATH)) model.to(device)如果保存模型在 CPU 上，加载模型在 GPU 上，那么我们的保存与加载的代码应该如下：
torch.save(model.state_dict(), PATH) device = torch.device("cuda") model = Model(*args, **kwargs) # 选择加载到哪一个GPU设备上 model.load_state_dict(torch.load(PATH, map_location="cuda:0")) model.to(device)如果保存模型在 CPU 上，加载模型在 CPU 上，那么我们的保存与加载的代码应该如下：
torch.save(model.state_dict(), PATH) model = Model(*args, **kwargs) model.load_state_dict(torch.load(PATH))