I get this error when I was using pytorch: RuntimeError: gather_out_cuda(): Expected dtype int64 for index

Question

I'm trying to reproduce meshed memory transformer https://github.com/aimagelab/meshed-memory-transformer, but I get this error: RuntimeError: gather_out_cuda(): Expected dtype int64 for index. The detailed error information are as follows:

Traceback (most recent call last):
  File "/home/ai/data/meshed-memory-transformer/train.py", line 252, in <module>
    scores = evaluate_metrics(model, dict_dataloader_val, text_field)
  File "/home/ai/data/meshed-memory-transformer/train.py", line 55, in evaluate_metrics
    out, _ = model.beam_search(images, 20, text_field.vocab.stoi['<eos>'], 5, out_size=1)
  File "/home/ai/data/meshed-memory-transformer/models/captioning_model.py", line 70, in beam_search
    return bs.apply(visual, out_size, return_probs, **kwargs)
  File "/home/ai/data/meshed-memory-transformer/models/beam_search/beam_search.py", line 82, in apply
    visual, outputs = self.iter(t, visual, outputs, return_probs, **kwargs)
  File "/home/ai/data/meshed-memory-transformer/models/beam_search/beam_search.py", line 132, in iter
    self.model.apply_to_states(self._expand_state(selected_beam, cur_beam_size))
  File "/home/ai/data/meshed-memory-transformer/models/containers.py", line 30, in apply_to_states
    self._buffers[name] = fn(self._buffers[name])
  File "/home/ai/data/meshed-memory-transformer/models/beam_search/beam_search.py", line 38, in fn
    beam.expand(*([self.b_s, self.beam_size] + shape[1:])))
RuntimeError: gather_out_cuda(): Expected dtype int64 for index

GPU version:RTX3090

Cuda compilation tools, release 11.2, V11.2.142
Build cuda_11.2.r11.2/compiler.29558016_0

PyTorch version:1.9.1   build:py3.7_cuda11.1_cudnn8.0.5_0