RuntimeError: expected device cpu but got device cuda:0

最新推荐文章于 2024-12-31 12:47:32 发布

原创最新推荐文章于 2024-12-31 12:47:32 发布 · 1.7k 阅读

3 ·

CC 4.0 BY-SA版权

pytorch 专栏收录该内容

28 篇文章

订阅专栏

在使用PyTorch进行模型微调训练时遇到RuntimeError，原因是尝试在CPU上操作一个位于CUDA设备上的张量。解决办法是确保在加载优化器状态时，将所有张量转移到CUDA设备上。示例代码展示了如何遍历优化器的状态字典，将张量转换为CUDA设备上的张量。

在对模型进行微调训练的时候，代码报错如下：

File "/home/xxxx/EdgeNets/utilities/train_eval_seg.py", line 48, in train_seg
    optimizer.step()
  File "/home/xxxx/anaconda3/envs/pytorch/lib/python3.6/site-packages/torch/optim/sgd.py", line 100, in step
    buf.mul_(momentum).add_(1 - dampening, d_p)
RuntimeError: expected device cpu but got device cuda:0

解决的方法是，重载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()