Beam search error about meshed-memory-transformer

[Faiail]

Actually I checked that during beam search the shape of previous generated token changes from bs_ X 1 to 40 X 1 and I trurly don't know why. In particular this happens after generating the first token. In fact the first input is a batch of start tokens, with the correct batch size, whereas after the first iteration the it tensor changes the shape. This is related to the fact that it is taken the previous generated token as input. How can I fix that?

[KillerGary]

RuntimeError: gather(): Expected dtype int64 for index

[Kuo2023]

RuntimeError: gather(): Expected dtype int64 for index:
I get this too, I have to add .to(torch.int64) at the end of every line of the code with gather (e.g
s = torch.gather(s.view(([self.b_s, cur_beam_size] + shape[1:])), 1,
beam.expand(([self.b_s, self.beam_size] + shape[1:])).to(torch.int64))
)
And then I am encountered with other different problems
RuntimeError: Expected tensor for argument #1 'indices' to have one of the following scalar types: Long, Int; but got torch.cuda.FloatTensor instead (while checking arguments for embedding)