【pytorch】模型训练和推理使用的设备不一致导致的问题

问题

前几天使用0号显卡训练了模型，今天想要推理的时候，0卡被占用，所以选择使用1卡进行推理。
但是推理的时候却报错：RuntimeError: CUDA error: out of memory

原因

所有先前保存的模块，不论其设备，首先加载到CPU上，然后移动到之前保存它们的设备上（也就是0号显卡）。因为0号显卡被占用，所以显示out of memory

解决方法

在load模型的时候，添加map_location参数

torch.load(path_file, map_location='cuda:1')

参考资料

1. https://github.com/apachecn/pytorch-doc-zh/blob/master/docs/1.0/jit.md
2. pytorch在gpu训练加载模型,训练与测试在不同机器上遇到的问题总结
3. pytorch将cpu训练好的模型参数load到gpu上，或者gpu->cpu上