PyTorch DataParallel 指定显卡运行出错

发表于 更新于

问题描述：当设置的显卡ID中包含第0张显卡时，程序能够正常运行；若指定从显卡1开始，例如我想要使用显卡1和显卡2并行计算，此时会出现如下错误：

1
RuntimeError: Expected tensor for argument #1 'input' to have the same device as tensor for argument #2 'weight'; but device 1 does not equal 0 (while checking arguments for cudnn_batch_norm)

原因：我们的数据默认存在第0张显卡，但是选用的显卡里面并不包含第0张显卡，因此就会出错。
解决方法：在代码中加入

1
torch.cuda.set_device(1)

即可改变默认权重默认的存储位置，程序就能够正常运行。