最近在学习使用pytorch，遇到一个坑。
从中断的训练中继续训练需要重新加载检查点，一些优化器（比如adam）的一些变量也需要被保存到检查点中。
而在使用load_state_dict()还原时，有可能会将这些变量还原到CPU上，如下面的代码这样：
model = Model()
optimizer = Optim()
checkpoint = torch.load(ckp_path)
model.load_state_dict(checkpoint["model"])
optimizer.load_state_dict(checkpoint["optimizer"])
model.cuda().train()
......像这样的话，优化器的变量就会被还原到CPU上，这在训练更新权重时会导致错误，因为模型权重在GPU上，而优化器的变量在CPU上。
造成这个问题的原因是在还原检查点之前，模型没有被转到GPU上，这样还原时，所有模型相关的变量都会被还原到CPU上。
解决办法有三。
方法一：手动转换
model = Model()
optimizer = Optim()
checkpoint = torch.load(ckp_path)
model.load_state_dict(checkpoint["model"])
optimizer.load_state_dict(checkpoint["optimizer"])
for state in optimizer.state.values():
for k, v in state.items():
if isinstance(v, torch.Tensor):
state[k] = v.cuda()
model.cuda().train()
......方法二：还原检查点之前，将模型转到GPU
model = Model().cuda()
optimizer = Optim()
checkpoint = torch.load(ckp_path)
model.load_state_dict(checkpoint["model"])
optimizer.load_state_dict(checkpoint["optimizer"])
model.cuda().train()
......方法三：在load中设置map_location
device = torch.device("cuda")
model = Model()
optimizer = Optim()
checkpoint = torch.load(ckp_path, map_location=device)
model.load_state_dict(checkpoint["model"])
optimizer.load_state_dict(checkpoint["optimizer"])
model.cuda().train()
......