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?