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()
            
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值