复现Meshed Memory Transformer时PyTorch报gather_out_cuda索引需int64错误

问题根因 #

该报错是因为PyTorch 1.9+版本的CUDA gather算子要求输入的索引张量必须为int64（long）类型，而该项目原始代码中beam搜索阶段生成的选中束索引selected_beam为int32类型，不符合算子要求。

解决方案 #

• 第一步：定位到项目路径下的models/beam_search/beam_search.py文件
• 第二步：修改_expand_state方法的返回值，将选中的束索引强制转换为long类型，找到方法中返回selected_beam的代码行，修改为return selected_beam.long()
• 第三步（如果修改后仍报错可补充修改）：找到iter方法中调用_expand_state的代码行，提前做类型转换：
  将原行

  self.model.apply_to_states(self._expand_state(selected_beam, cur_beam_size))
  

  修改为

  self.model.apply_to_states(self._expand_state(selected_beam.long(), cur_beam_size))
  

• 第四步：保存修改后重新运行评估流程即可解决该报错

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