lets say I have tensor of shape (2, 10) .
I want to create a mask tensor with same shape that is ‘1’ only immediately after occurrence of a specific value and ‘0’ elsewhere.
For example considering I want to mask indices after 3:

``````x = tensor([[1, 3, 4, 5, 3, 6],
[1, 2, 5, 3, 7, 9]])
mask = tensor([[False,  False, True, False,  False, True],
[False, False, False,  False, True, False]])
``````

x.eq(3) will mask indices of 3 but not the next index.

Thanks for ur help

Just x== value, you can get the result you want

I want the next element of the index with that specific value be ‘True’ not the index of value

You can return the index of value by torch.nonzero(x==3), then manipulate the index, to generate the mask.

Yeah actually I tried that but I could not generate what I want since each row may have multiple occurrence of that value. I don’t know how to manipulate the result of torch.nonzero(x==3).

``````ind = (x==3).nonzero()
tensor([[0, 1],
[0, 4],
[1, 3]])
``````

Should I do it in a for loop? Or is there a nice vectorized way of doing it?

===== 1

``````mask = x == 3
``````

===== 2

``````row, col = torch.where(x == 3)
col += 1
``````mask = torch.zeros_like(x).bool()