Solved by justusschock in post #2 
              
              
                You could use binary masking to achieve this: 
seq_lens = torch.tensor([2,3,5,1]).unsqueeze(-1)
max_len = torch.max(seq_lens)

# create tensor of suitable shape and same number of dimensions
range_tensor = torch.arange(max_len).unsqueeze(0)
range_tensor = range_tensor.expand(seq_lens.size(0), range_…