在进行跨平台模型加载时遇到报错ModuleNotFoundError: No module name 'torch_npu'

报错信息

• 报错截图

  

• 报错文本

  ......
  ModuleNotFoundError: No module name 'torch_npu'

问题分析

在进行跨平台模型加载时，调用torch.load()，遇到报错“ModuleNotFoundError: No module name 'torch_npu' ”。这是由于在保存模型时，没有将模型转换为CPU模型或者保存了NPU optimizer的信息。

处理方法

在NPU环境中执行代码，加载模型权重文件中的参数。后续再把模型放到CPU上再进行存储，存储后的文件将不依赖torch_npu包进行解析加载。样例解决方法如下所示。

import torch
import torch_npu
...
tmp = torch.load('xxx.pt',map_location='cpu')
torch.save(tmp, 'new_xxx.pt')