optimizer.step()处报错：RuntimeError: Expected all tensors to be on the same device, but found cuda:0

1、pytorch在加载预训练模型续训时在optimizer.step()处报错：RuntimeError: Expected all tensors to be on the same device, but found cuda:0 and cpu!。

主要原因是：optimizer加载参数时，tensor默认在CPU上，当你使用GPU训练时就会报以上错误。

2、解决办法，在optimizer加载完全部参数后将optimizer里的tensor数据全部转到GPU上，代码如下：

checkpoint = torch.load("xxxx.pth"),map_location=lambda storage, loc: storage)   # Load all tensors onto the CPU
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['opt_dict'])
# 增加以下几行代码，将optimizer里的tensor数据全部转到GPU上
for state in optimizer.state.values():
    for k, v in state.items():
        if torch.is_tensor(v):
            state[k] = v.cuda()