I’m trying to create *_key_padding_mask for torch.nn.Transformer. So I want to create a mask using the length of sentence data in the batch.
Assuming number of time steps  t=7 t=8
Edit: changed value of t
batch_size = 5
idx = torch.tensor([3,4,2,1,6])


mask = torch.tensor([[0,0,0,1,1,1,1,1],
                     [0,0,0,0,1,1,1,1],
                     [0,0,1,1,1,1,1,1],
                     [0,1,1,1,1,1,1,1],
                     [0,0,0,0,0,0,1,1]])

How do I get mask using idx with vectorized code?