I have a tensor
JavaScript
x
2
1
t = torch.tensor([[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [1, 0, 0, 0]])
2
and a query tensor
JavaScript
1
2
1
q = torch.tensor([1, 0, 0, 0])
2
Is there a way to get the indexes of q
like
JavaScript
1
2
1
indexes = t.index(q) # get back [0, 3]
2
in pytorch?
Advertisement
Answer
How about
JavaScript
1
5
1
In [1]: torch.nonzero((t == q).sum(dim=1) == t.size(1))
2
Out[1]:
3
tensor([[ 0],
4
[ 3]])
5
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:
JavaScript
1
2
1
torch.all(t==q, dim=1)
2