I adapted my code in the ways you suggested and now it works.
Thank you again.
I modified the masking function, based on another thread so that it now looks like this:
def create_seq_mask(seq_lengths : torch.tensor, device):
    length = int(seq_lengths.max())
    res = torch.stack([torch.arange(length, device=device) < x for x in seq_lengths])
    return res.unsqueeze(2)