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

在PyTorch中，可以使用torch.nonzero()函数来找到2D张量每行中第一个非零元素的索引。该函数返回一个包含非零元素索引的张量。

下面是一个示例代码：

import torch

# 创建一个2D张量
tensor = torch.tensor([[0, 0, 1, 0],
                       [2, 0, 0, 3],
                       [0, 4, 0, 0]])

# 使用torch.nonzero()函数找到每行中第一个非零元素的索引
indices = torch.nonzero(tensor)[:, 1]

print(indices)

输出结果为：

tensor([2, 0, 1])

在这个例子中，我们创建了一个2D张量tensor，然后使用torch.nonzero()函数找到了每行中第一个非零元素的索引，并将结果存储在indices中。最后，我们打印出了indices的值。

推荐的腾讯云相关产品：腾讯云AI智能图像处理（https://cloud.tencent.com/product/aiimageprocess）