Hi Roxor!



 roxor:

I am looking for a possibility to get ALL indices for the maximum for each row in a 2-D tensor.
Let us define the following tensor:
a = torch.tensor([[5,8,9,6,9], [5,4,3,2,1]])



Try this:
>>> torch.__version__
'1.7.1'
>>> a = torch.tensor([[5,8,9,6,9], [5,4,3,2,1]])
>>> torch.nonzero ((a == a.max (dim = 1, keepdim = True)[0]))
tensor([[0, 2],
        [0, 4],
        [1, 0]])



torch.tensor([[2,4],[0]])



Note, that pytorch does not support “ragged tensors” (tensors whose rows
are not equal to one another in length), so you can’t get your result in a
format like this.
Best.
K. Frank