淮南创业网/优化内容

问题：

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

现象：

• 模型训练是在 cuda:1 上训练
• 在预测时指定 cuda:0 不work，出现上述问题

解决方案：

• 首先在 predict.py 中修改：

  #device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu") device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

  • 这样的修改将导致上述问题，
  • 因为model在 cuda:0 中，但是具体的model处理数据的流程在 cuda:1 中。（ model_structure.py 的 class TwoTowerModel ）

• 具体解决方法：将 model_structure.py 的 class TwoTowerModel(nn.Module): 中的 forward 函数中增加一行

  self.device = torch.device("cuda:0")