import torch

batch_size = 5

idx = torch.tensor([3,4,1,1,6])

# I changed the mask a little

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]])

print(torch.gather(mask, dim=1, index=idx.unsqueeze(-1).expand(mask.size(0), -1)))