pytorch加载模型继续训练时optimizer参数被加载到cpu

问题

在使用pytorch加载模型继续进行训练时，报错’ Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!’

报错意思就是所有tensors应该被加载至同一位置，但有些tensor在CPU有些在GPU

解决

首先torch.load(load_file)在加载模型时默认加载到CPU上，但我希望加载到GPU上继续训练，于是使用’model.to(device)'将模型参数加载到GPU，仍然报错。看到代码里还加载了一个optimizer，那么问题一定就是optimizer优化器里的参数没有被加载到GPU，输入以下内容完美解决（其实抓耳挠腮地花了一晚上才解决。。。）：

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

感谢

优化器导入参数导致数据出现在cpu和gpu上