Pytorch使用训练好的模型时报错RuntimeError: cuda runtime error (10) : invalid device ordinal

跑了两天的模型，拿出来做预测的时候加载出了问题，报错如下：

RuntimeError: CUDA error (10): invalid device ordinal (check_status at /opt/conda/conda-bld/pytorch_1532579805626/work/aten/src/ATen/cuda/detail/CUDAHooks.cpp:36)

最后发现是因为自己训练的时候用了多卡：

        就是训练代码里的这两行代码惹的祸

os.environ['CUDA_VISIBLE_DEVICES'] = "0, 1, 2, 3" net_model = nn.DataParallel(net_model, device_ids=[0, 1, 2, 3])

 

针对这种情况，解决方法如下：

        将预测代码里的这两行改为如下：

os.environ['CUDA_VISIBLE_DEVICES'] = "0, 1, 2, 3" net_model = nn.DataParallel(net_model, device_ids=[0])