Pytorch Error:RuntimeError:CUDA error: out of memory [in torch.load]

Error:

当加载模型时：

model.load_state_dict(torch.load('xxx.pth'))

Error:RuntimeError:CUDA error: out of memory

解决方法

原因1-CUDA显存不足

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

原因2-显卡被占用

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

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

原因3-模型保存GPU不同(系统默认map到0卡)

使用cuda：0训练保存的模型，如果直接使用cuda：1去load，也会报错
解决办法：model.load_state_dict(torch.load(‘trained.pth’, map_location={‘cuda:0’: ‘cuda:1’}))
或者
用同一块显卡去Load

除此之外，test时加上：

with torch.no_grad():
  # test process

原因3的根本问题应该在于torch.load，不应该是保存的GPU不同所带来的，十分感谢评论区的指点，链接里面有比较详细的介绍，可以参考