How to get indices of multiple elements in a 2D tensor, in a GPU friendly way?
This question is similar to that already answered here, but that question does not address how to retrieve the indices of multiple elements.

I have a 2D tensor points with many rows and a small number of columns, and would like to get a tensor containing the row indices of all the elements in that tensor. I know what elements are present in points beforehand; It contains integer elements ranging from 0 to 999, and I can make a tensor using the range function to reflect the set of possible elements. The elements may be in any of the columns.

How can I retrieve the row indices where each element appears in my tensor in a way that avoids looping or using numpy, so I can do this quickly on a GPU?

I am looking for something like (points == elements).nonzero()[:,1]

Thanks!

try torch.cat([(t == i).nonzero() for i in elements_to_compare])

&gt;&gt;&gt; import torch
&gt;&gt;&gt; t = torch.empty((15,4)).random_(0, 999)
&gt;&gt;&gt; t
tensor([[429., 833., 393., 828.],
        [555., 893., 846., 909.],
        [ 11., 861., 586., 222.],
        [232.,  92., 576., 452.],
        [171., 341., 851., 953.],
        [ 94.,  46., 130., 413.],
        [243., 251., 545., 331.],
        [620.,  29., 194., 176.],
        [303., 905., 771., 149.],
        [482., 225.,   7., 315.],
        [ 44., 547., 206., 299.],
        [695.,   7., 645., 385.],
        [225., 898., 677., 693.],
        [746.,  21., 505., 875.],
        [591., 254.,  84., 888.]])


&gt;&gt;&gt; torch.cat([(t == i).nonzero() for i in [7,385]])
tensor([[ 9,  2],
        [11,  1],
        [11,  3]])

&gt;&gt;&gt; torch.cat([(t == i).nonzero()[:,1] for i in [7,385]])
tensor([2, 1, 3])


Numpy:

&gt;&gt;&gt; np.nonzero(np.isin(t, [7,385]))
(array([ 9, 11, 11], dtype=int64), array([2, 1, 3], dtype=int64))

&gt;&gt;&gt; np.nonzero(np.isin(t, [7,385]))[1]
array([2, 1, 3], dtype=int64)


