python - Pytorch:如何找到 2D 张量的每一行中第一个非零元素的索引？

Question

我有一个二维张量，每行都有一些非零元素，如下所示:

import torch
tmp = torch.tensor([[0, 0, 1, 0, 1, 0, 0],
                    [0, 0, 0, 1, 1, 0, 0]], dtype=torch.float)

我想要一个包含每行中第一个非零元素索引的张量:

indices = tensor([2],
                 [3])

如何在 Pytorch 中计算它？

Answer 1

我简化了 Iman 的方法来执行以下操作:

idx = torch.arange(tmp.shape[1], 0, -1)
tmp2= tmp * idx
indices = torch.argmax(tmp2, 1, keepdim=True)