sten
sten

Reputation: 7486

Search a tensor for data/values

Given:

tensor([[6, 6],
        [4, 8],
        [7, 5],
        [7, 4],
        [6, 4]])

How do I find the index of rows with values [7,5]?

In general, how do I search for indices of any values: full and partial row or column?

Upvotes: 0

Views: 126

Answers (1)

U13-Forward
U13-Forward

Reputation: 71580

Try with this:

>>> (a[:, None] == torch.tensor([7, 5])).all(-1).any(-1).nonzero().flatten().item()
2
>>> 

Upvotes: 1

Related Questions