从张量tensorflow中获取索引

Question

all_max = tf.convert_to_tensor([[4, 2, 3], [3, 4, 5]], dtype=tf.float32)

如何从张量数组all_max中得到元素[3,4,5]的索引?

在list中，我们简单地使用list.index(element)来获取list中存在的元素的索引。

感谢

Answer 1

我发现这个在:https://www.py4u.net/discuss/147615

要找到2d/3d张量的元素索引，将其转换为1d# ie的例子。查看(元素个数)

的例子:

mat=torch.tensor([[1,2],[4,3])
#to find index of 2
five = 2
mat=mat.view(4)
numb_of_col = 4
for o in range(numb_of_col):
if mat[o] == five:
print(torch.tensor([o]))