当前位置: 首页 > 工具软件 > Location Map > 使用案例 >

please use torch.load with map_location=torch.device(‘cpu‘) to map your storages to the CPU.

宇文嘉勋
2023-12-01

please use torch.load with map_location=torch.device(‘cpu’) to map your storages to the CPU.

问题描述:

在GPU环境下,用pytorch框架训练好的网络模型,在CPU环境测试,报错如下:

RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False. If you are running on a CPU-only machine, please use torch.load with map_location=torch.device('cpu') to map your storages to the CPU.

RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False. If you are running on a CPU-only machine, please use torch.load with map_location=torch.device(‘cpu’) to map your storages to the CPU.

可定位到报错行的代码是模型加载部分:

torch.load('xxxNet.pth')

其中,‘xxxNet.pth’ 是用GPU训练好的模型权重文件。


解决方案:

在模型加载时,指定 map_location='cpu' 即可,代码如下:

torch.load('xxxNet.pth', map_location='cpu')
 类似资料: