pytorch加载模型torch.load报cuda out of memory解决方案

当服务器中有多张显卡时可能会出现这个问题。
模型参数加载：

model_recover = torch.load(args.model_recover_path)

报错如下：

RuntimeError: CUDA error: out of memory

原因：
当使用torch.load加载模型参数时，会默认加载在第一块GPU0上，当GPU0没有在使用时，问题不大，但是显存被占满时这时候就加载不了了。
解决方法：

model_recover = torch.load(args.model_recover_path, map_location={
   
   'cuda:0': 'cuda:2'})

指定GPU映射，将GPU0映射到GPU2（任意一张空闲显卡），这样原本默认加载在GPU0上的参数将会加载到GPU2中，问题解决。

参考：torch代码运行时显存溢出问题之问题四

来源：oschina

链接：https://my.oschina.net/u/4255339/blog/4706560