pytorch重载optimizer参数时报错：RuntimeError: expected device cpu but got device cuda:0的解决方法
·
1. 问题描述：
我在使用torch.save()保存了optimizer的参数过后，
torch.save(
{
'state_dict':net.state_dict(),
'optimizer':optimizer.state_dict(),
'epochID':epoch,
},
filename
)
再次利用optimizer.load_state_dict()加载参数，在optimizer.step()处报错：
RuntimeError: expected device cpu but got device cuda:0
2. 解决办法：
重载optimizer的参数时将所有的tensor都放到cuda上（加载时默认放在cpu上了），代码片段如下：
checkpoint = torch.load(filename)
net.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
for state in optimizer.state.values():
for k, v in state.items():
if torch.is_tensor(v):
state[k] = v.cuda()
current_epoch = checkpoint['epochID'] + 1
顺利解决问题～
北京人形旗下天工造物具身智能开源社区，聚焦具身天工与慧思开物两大平台
更多推荐
18
0- 0
已为社区贡献5条内容
所有评论(0)