Pytorch Tensor 如何获取元素的索引？ [复制]

Question

我有 2 个名为 x 和 list 的张量，它们的定义如下：

x = torch.tensor(3)
list = torch.tensor([1,2,3,4,5])

现在我想从 list 中获取元素 x 的索引。预期的输出是一个整数：

2

我怎样才能轻松完成？

Answer 1

import torch

x = torch.tensor(3)

list = torch.tensor([1,2,3,4,5])
idx = (list == x).nonzero().flatten()
print (idx.tolist()) # [2]

list = torch.tensor([1,2,3,3,5])
idx = (list == x).nonzero().flatten()
print (idx.tolist()) # [2, 3]