torch.load CUDA ERROR: out of memory

在pytorch中load模型，出现CUDA ERROR:out of memory,基本有下面三个原因：

CUDA显存不足

解决办法：换成显存更大的显卡

显卡被占用

当存在多个显卡时，pytorch默认使用0号显卡，当0号显卡被占用时，即使1号显卡空闲仍然会报错
解决办法：让其他显卡不可见，此时0号显卡为空余的1号显卡

import os

os.environ['CUDA_VISIBLE_DEVICE']='1'

模型保存GPU不同

使用cuda：0训练保存的模型，如果直接使用cuda：1去load，也会报错
解决办法：model.load_state_dict(torch.load('trained.pth', map_location={'cuda:0': 'cuda:1'}))

参考