Pytorch Tensor 如何获取元素的索引? [英] How Pytorch Tensor get the index of elements?

问题描述

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

I have 2 Tensors named x and list and their definitions are below:

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

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

Now I want to get the index of element x from list. The expected output is an Integer:

2

我怎样才能以简单的方式进行?

How can I do in an easy way?

推荐答案

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]

这篇关于Pytorch Tensor 如何获取元素的索引?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！