Machine Learning-CUDA out of memory (torch.load())解决

关于Pytorch, Tensorflow等等的环境配置，可以查看这个专栏哦

环境配置(Pytorch, Tensorflow, Linux, etc.)

Machine Learning-CUDA out of memory 解决

杂话

在上次的CUDA out of memory的问题解决之后（CUDA out of memory (Training)解决），这次又是遇到了新的问题。
使用Pytorch的朋友可能都知道，我们是可以使用torch.load()来加载我们之前训练好的模型的。

问题

于是今天，我在使用torch.load()的时候就又碰到了爆显存的问题

但是我仔细看了一下我的代码，我在加载模型的时候，好像也没有用到那块被占用的显卡啊
（自从上次的报错之后，我和实验室的另一位同学达成了一人使用一块显卡的和平相处协议）

感谢这位大哥的帖子，帮我解决了问题pytorch加载模型torch.load报cuda out of memory解决方案

解决

简单来说就是，在torch.save()得到的pth文件里，会保存着你之前是在那张显卡训练这个模型的信息，而你在加载这个模型的时候，就会经过这张显卡，这个时候，哪怕你最终不是使用这张显卡计算，但是这张显卡现在是被占用的话，就会报CUDA out of memory的错误。。。

也就是说，因为我和大哥一人一张显卡，之前我一直都是在cuda:1进行训练的，但是今天1卡被大哥用了，就算我把device改为cuda:0，在加载模型的时候也还是会报错，这个时候你就需要在代码中加入一个参数，来直接把模型映射到你要使用的显卡上

于是代码就从
torch.load(resume_path)
变成了
torch.load(resume_path, map_location = 'cuda:0')

这样问题就解决咯

注

如果还有什么报错的话，也欢迎大家在评论区继续讨论哦！