How to set all elements of pytorch tensor to zero after a certain index in the given axis, where the index is given by another pytorch tensor?
I've two PyTorch tensors
mask = torch.ones(1024, 64, dtype=torch.float32)
indices = torch.randint(0, 64, (1024, ))

For every ith row in mask, I want to set all the elements after the index specified by ith element of indices to zero. For example, if the first element of indices is 50, then I want to set mask[0, 50:]=0. Is it possible to achieve this without using for loop?
Solution with for loop:
for i in range(mask.shape[0]):
    mask[i, indices[i]:] = 0


You can first generate a tensor of size (1024x64) where each row has numbers arranged from 0 to 63. Then apply a logical operation using the indices reshaped as (1024x1)
mask = torch.ones(1024, 64, dtype=torch.float32)
indices = torch.randint(0, 64, (1024, 1))    # Note the dimensions

mask[torch.arange(0, 64, dtype=torch.float32).repeat(1024,1) &gt;= indices] = 0

