我在使用 Torch 加载模型时遇到问题,该模型是使用 GPU 训练的,我正在尝试使用 CPU 加载该模型。然而,我成功地加载了模型,但在预测结果时出现错误。但是,如果我使用 GPU 机器,我可以预测输出,但不能预测 CPU 上的输出:
我的代码:
****To save the model I am using :****
PATH = "model.pt"
torch.save(model, PATH)
**To Load the Model**
import torch
PATH = "model.pt"
device = torch.device('cpu')
loaded_model=torch.load(PATH, map_location=device)
我能够成功加载模型。但在预测时我遇到了运行时错误
**Predicting the loaded model using CPU**
predicted_title = loaded_model.predict([abstract])
运行时错误:在您的系统上找不到 NVIDIA 驱动程序。请检查您是否有 NVIDIA GPU 并安装了驱动程序
如果错误可能非常简单,我很抱歉,但我无法纠正这个问题。
您可以使用
输出型号设备print(loaded_model.device)
如果不是CPU,则
model = model.to('cpu')