PyTorch DataParallel 指定显卡运行出错

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

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张显卡，因此就会出错。
解决方法：
加入如下语句

torch.cuda.set_device(1)  

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