GPU训练的PyTorch模型加载到CPU预测时报无NVIDIA驱动错误如何解决

Runtime Error: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver

问题根因 #

报错触发在预测阶段，与模型加载时的map_location配置无关，核心原因是推理逻辑中仍存在调用CUDA设备的硬编码逻辑，常见场景如下：

• 自定义的predict方法内部固定将输入张量、中间计算张量移动到cuda设备，例如写死input = input.to("cuda")
• 你使用了保存整个模型对象的方式（而非仅保存模型参数state_dict），模型内部绑定的训练阶段CUDA设备属性没有被map_location完全覆盖，部分层的计算逻辑仍调用原CUDA设备

修复方案 #

方案1：修复predict方法设备逻辑（推荐，可一次性兼容CPU/GPU环境） #

修改predict方法中设备绑定的硬编码逻辑，自动适配模型当前所在的设备，参考修改代码：

def predict(self, input_data):
    # 获取模型当前所在的设备，替代固定写死的cuda
    device = next(self.parameters()).device
    # 输入处理完成后移动到对应设备
    input_tensor = self.preprocess(input_data).to(device)
    # 剩余推理逻辑保持不变
    with torch.no_grad():
        output = self(input_tensor)
    return self.postprocess(output)

修改后重新保存模型，后续加载无需额外调整即可同时兼容CPU、GPU环境。

方案2：无需重新保存模型的临时修复 #

加载模型后手动完成全模型的CPU迁移，调整加载和预测代码如下：

import torch
PATH = "model.pt"
device = torch.device('cpu')
loaded_model=torch.load(PATH, map_location=device)
# 新增：手动触发全模型设备迁移
loaded_model = loaded_model.to(device)
loaded_model.eval()
# 若predict方法仍存在输入硬编码到cuda的逻辑，可提前将输入处理为CPU张量再传入
# 示例：
# processed_input = preprocess([abstract]).to(device)
# predicted_title = loaded_model.predict(processed_input)
predicted_title = loaded_model.predict([abstract])

方案3：长期优化模型保存规范 #

建议后续使用仅保存模型参数state_dict的方式存储模型，避免设备绑定、PyTorch版本兼容、代码结构依赖等问题：

模型保存代码

PATH = "model_state.pt"
torch.save(model.state_dict(), PATH)

模型加载代码

import torch
# 从模型定义文件导入你的模型类
from your_model_def import YourModelClass
device = torch.device('cpu')
# 初始化模型结构
model = YourModelClass()
# 加载参数
model.load_state_dict(torch.load(PATH, map_location=device))
model = model.to(device)
model.eval()

内容的提问来源于stack exchange，提问作者Harry DSOUZA