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

记录一下遇到的改进网络模型时的错误。

起因：将CCNet的十字交叉注意力模块加入YOLOv5中，在写代码时，有部分方法用到了.cuda()这个方法来返回一个临时数据。YOLOv5在构建模型时，会跑一遍模型，此时数据类型是cpu类型的，当与cuda类型的临时数据进行计算时，自然会发生冲突。如果直接去掉.cuda()方法，临时数据变为cpu类型；虽然可以在第一次构建模型时不发生冲突，但是实际训练的时候，传入的数据是cuda类型的，因此会再次发生冲突。

解决方案：在传入的数据和临时数据计算之前，先将临时数据的类型转换为传入数据的类型就可以了。

def INF(B, H, W):
    return -torch.diag(torch.tensor(float("inf")).repeat(H), 0).unsqueeze(0).repeat(B * W, 1, 1)
    # return -torch.diag(torch.tensor(float("inf")).cuda().repeat(H), 0).unsqueeze(0).repeat(B * W, 1, 1)

......


energy_H = (torch.bmm(query_H, key_H) + self.INF(b2, h1, w1).to(query_H.device)).view(b2, w1, h1, h1).permute(0, 2, 1, 3)

如上面的代码所示，解决类型冲突的方法，就是令self.INF().to(传入数据.device)