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)