Issue with Multiple `eos_token_id` in Code Sequences

[qixuema]

Hi, Phil,

I've noticed a potential issue in the code where multiple eos_token_id are being added to the code sequences, instead of just one at the end of each row. It seems to insert four eos_token_id values. Could you please clarify if this is the expected behavior?

batch_arange = torch.arange(batch, device = device)
batch_arange = rearrange(batch_arange, '... -> ... 1')

codes[batch_arange, code_lens] = self.eos_token_id

Best regards, Xueqi

[lucidrains]

@qixuema yes indeed there was an issue of a missing dimension

thank you Xueqi!