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()