Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

问题描述

注意：任何标题提示的报错都可以用下面方法的原理找出问题数据（在CPU），然后将该数据放入GPU即可
在运行bert模型时，打算将代码放到服务器的GPU上面运行，将模型和设置的参数

if torch.cuda.is_available():
    device = torch.device('cuda')
else:
    device = torch.device('cpu')
gpus = [0, 1]

模型的搭建如下

model = BERT().to(device)
model = DataParallel(model, device_ids=[gpus], output_device=[gpus])

数据从cpu到gpu的迁移如下

input_ids = input_ids.to(device)
masked_tokens = masked_tokens.to(device)
masked_pos = masked_pos.to(device)

运行后报错如下

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking arugment for argument mat2 in method wrapper_mm)
可见是部分数据在cpu，部分数据在gpu导致的问题，但是将所有代码中能找到的tensor数据都用了xxx.to(device)，还是这样报错尝试了很多方法无果，后来发现并不是单纯的把某个数据.to(device）就可以解决的，因为我已经几乎把所有的数据都.to(device)了，后来通过在错误行设置断点，在该行进行单步调试，然后进入下一层函数后再通过错误行设断点，再单步调试最终发现是parameter不在cpu，然后我对某个函数中的parameter设置.to(device)后还是无法解决问题，最后发现是把网络层放在forward导致的，因为在forward搭建的网络不进入GPU。

解决方案

通过检查自己定义的所有类，发现这个类有问题

经过检查发现def forward内部定义了两个网络层，如上图红圈处，将这两个层定义到__init__（）中，然后在forward调用这两个层，修改后如下图

问题解决！！！