Pytorch在forward中创建了新的tensor时，RuntimeError: arguments are located on different GPUs

最新推荐文章于 2021-06-02 22:52:57 发布

(ーー゛)

最新推荐文章于 2021-06-02 22:52:57 发布

阅读量798

点赞数

文章标签： cuda gpu

本文链接：https://blog.csdn.net/weixin_44430537/article/details/107412415

版权

Pytorch在forward中创建了新的tensor时，RuntimeError: arguments are located on different GPUs

这个问题浪费了我一天时间，找了很多关于相关错误提示的解决方案。很不幸，有用但解决不了。
下面进入正题：
我的错误引起是：
我在自定义forward中声明了一个tensor

class test(nn.Module):
    def __init__(self, xxxx):
        super(test, self).__init__()
        .....
    def forward(self, x):
        outf = torch.zeros((2, 2, 2, 1))
        if torch.cuda.is_available():
            outf = outf.cuda()
        print(outf[0,0,0,0], x[0,0,0,0])
        outf = outf + x
        return outf

大概形式是这样的，然后报错。
原因：如截图，两个tensor用的GPU不同。

我看了下训练时采用是to(device)，当然device是用的4号GPU。当时我也想到会存在这种bug，然后又加了
os.environ[‘CUDA_VISIBLE_DEVICES’] = str(4)
但是竟然没用，可能是之前误以为加了这句会是.cuda()，默认使用设置那个GPU。