PyTorch torch.nonzero 函数
Pytorch torch 参考手册torch.nonzero 是 PyTorch 中用于返回非零元素索引的函数。
函数定义
torch.nonzero(input)
使用示例
实例
import torch
x = torch.tensor([[0, 1], [2, 0], [0, 3]])
# 返回非零元素的索引
indices = torch.nonzero(x)
print("非零元素索引:")
print(indices)
x = torch.tensor([[0, 1], [2, 0], [0, 3]])
# 返回非零元素的索引
indices = torch.nonzero(x)
print("非零元素索引:")
print(indices)
输出结果为：
非零元素索引:
tensor([[0, 1],
[1, 0],
[2, 1]])