【错误处理】RuntimeError: Attempting to deserialize object on CUDA device 1 but torch.cuda.device_count

1 项目场景

使用LSTM进行交叉路口的车辆行为识别。

---

2 问题描述

在使用PyTorch加载模型时出现错误语句：

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

---

3 原因分析

模型原本是使用两个GPU进行训练的，但是我的电脑仅有一个GPU，所以会报错。

---

4 解决方案

将

model = torch.load(model_path)

修改为：

model = torch.load(model_path, map_location='cuda:0')

如果是2块GPU：就把map_location改为：map_location={'cuda:1': 'cuda:0'}
其余情况以此类推。