解决：Expected tensor for argument #1 'input' to have the same device as tensor for argument #2 'weight

柯西一施瓦兹不等式

于 2019-12-19 15:52:28 发布

阅读量9k

点赞数 4

CC 4.0 BY-SA版权

文章标签： python 机器学习深度学习

本文链接：https://blog.youkuaiyun.com/weixin_38654615/article/details/103615806

本文解决PyTorch1.0环境下，使用pretrainedmodels库中的resnet50模型进行多GPU训练时出现的weight和input不在同一device上的错误。通过调整模型结构，将部分层指定到特定的GPU上，成功实现多GPU并行训练。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

问题描述：

pytorch1.0，从pretrainedmodels 中import 了 resnet50模型
单GPU训练没有问题
但是改为多GPU训练的时候会报weight 和input不在一个device上的错误

解决方案
加入如下代码就可以了呢！！：

class Net(nn.Module):
  def __init__(self, model):
    super(Net, self).__init__()
    self.l1 = nn.Sequential(*list(model.children())[:-1]).to('cuda:0')
    self.last = list(model.children())[-1]

  def forward(self, x):
    x = self.l1(x)
    x = x.view(x.size()[0], -1)
    x = self.last(x)
    return x

在这里插入图片描述
链接：https://github.com/Cadene/pretrained-models.pytorch/issues/120