RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 2.

最新推荐文章于 2024-03-19 21:47:38 发布

哎呦不错的温jay

最新推荐文章于 2024-03-19 21:47:38 发布

阅读量1.1w

点赞数 11

分类专栏： pytorch 文章标签： pytorch 人工智能 python

本文链接：https://blog.csdn.net/wenqiwenqi123/article/details/122258589

版权

pytorch 专栏收录该内容

4 篇文章 0 订阅

订阅专栏

本文介绍了如何解决在使用PyTorch加载模型时遇到的RuntimeError，原因在于训练时使用的GPU设备与运行时设备不一致。提供了解决方案，即通过map_location参数映射模型存储到正确设备上。

摘要由CSDN通过智能技术生成

解决报错RuntimeError: Attempting to deserialize object on CUDA device 2 but torch.cuda.device_count() is 2. Please use torch.load with map_location to map your storages to an existing device.

这个报错的原因是训练的时候用了编号为2的GPU，但是在你用pytorch load这个模型运行的时候使用的不是2号GPU。

解决方法很简单，在load的时候加一个map_location：

#model.load_state_dict(torch.load(model_path))  #会报错的写法
model.load_state_dict(torch.load(model_path,map_location={'cuda:2': 'cuda:0','cuda:1': 'cuda:0'}))  #将对应的gpu编号映射到正确的gpu上