使用enhuiz/vall-e遇torch.load返回dict导致AttributeError错误求助

问题重现 #

运行enhuiz/vall-e项目时出现如下错误：

File "vall-e/vall_e/main.py", line 30, in main ar = torch.load(args.ar_ckpt).to(args.device) AttributeError: 'dict' object has no attribute 'to'

该问题在项目相关讨论中被提及但未得到解决，可在安装了deepspeed==0.8.3的Colab环境复现。

修复步骤 #

• 问题根源：torch.load(args.ar_ckpt)加载得到的是模型的状态字典（dict类型），而非完整的模型实例，字典对象没有.to()方法。
• 修改main.py第30行代码：
  原代码：

  ar = torch.load(args.ar_ckpt).to(args.device)
  

  替换为（需先导入对应模型类，确保参数匹配）：

  from vall_e.models.ar import AR  # 根据项目结构导入正确的AR模型类
  ar_model = AR()  # 按模型要求传入初始化参数
  ar_model.load_state_dict(torch.load(args.ar_ckpt))
  ar = ar_model.to(args.device)
  

• 若checkpoint是多GPU训练生成的，需添加map_location保证加载兼容性：

  state_dict = torch.load(args.ar_ckpt, map_location=args.device)
  ar_model.load_state_dict(state_dict)
  ar = ar_model.to(args.device)
  

补充建议 #

优先使用项目内置的模型加载工具函数（如果有），避免手动处理状态字典时出现适配问题。

内容的提问来源于stack exchange，提问作者Jonathan Coletti