How Pytorch Tensor get the index of specific value
I think there is no direct translation from list.index()
to a pytorch function. However, you can achieve similar results using tensor==number
and then the nonzero()
function. For example:
t = torch.Tensor([1, 2, 3])
print ((t == 2).nonzero(as_tuple=True)[0])
This piece of code returns
1
[torch.LongTensor of size 1x1]
For multidimensional tensors you can do:
(tensor == target_value).nonzero(as_tuple=True)
The resulting tensor will be of shape number_of_matches x tensor_dimension
. For example, say tensor
is a 3 x 4
tensor (that means the dimension is 2), the result will be a 2D-tensor with the indexes for the matches in the rows.
tensor = torch.Tensor([[1, 2, 2, 7], [3, 1, 2, 4], [3, 1, 9, 4]])
(tensor == 2).nonzero(as_tuple=False)
>>> tensor([[0, 1],
[0, 2],
[1, 2]])
Can be done by converting to numpy as follows
import torch
x = torch.range(1,4)
print(x)
===> tensor([ 1., 2., 3., 4.])
nx = x.numpy()
np.where(nx == 3)[0][0]
===> 2