Skip to content
Advertisement

Pytorch tensor – How to get the indexes by a specific tensor

I have a tensor

t = torch.tensor([[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [1, 0, 0, 0]])

and a query tensor

q = torch.tensor([1, 0, 0, 0])

Is there a way to get the indexes of q like

indexes = t.index(q) # get back [0, 3]

in pytorch?

Advertisement

Answer

How about

In [1]: torch.nonzero((t == q).sum(dim=1) == t.size(1))
Out[1]: 
tensor([[ 0],
        [ 3]])

Comparing t == q performs element-wise comparison between t and q, since you are looking for entire row match, you need to .sum(dim=1) along the rows and see what row is a perfect match == t.size(1).


As of v0.4.1, torch.all() supports dim argument:

torch.all(t==q, dim=1)
User contributions licensed under: CC BY-SA
2 People found this is helpful
Advertisement