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?

user_3337089 0
Viewed 152

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
pytorch vectorization python
1 Answers

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) >= indices] = 0
user_150201110
Related