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])

 

你可能感兴趣的:(pytorch)