PyTorch torch.nonzero Function
PyTorch torch Reference Manual
torch.nonzerois a function in PyTorch used to return the indices of non-zero elements.
Function Definition
torch.nonzero(input)
Usage Examples
Example
import torch
x = torch.tensor([[0, 1], [2, 0], [0, 3]])
# Return the indices of non-zero elements
indices = torch.nonzero(x)
print("Indices of non-zero elements:")
print(indices)
x = torch.tensor([[0, 1], [2, 0], [0, 3]])
# Return the indices of non-zero elements
indices = torch.nonzero(x)
print("Indices of non-zero elements:")
print(indices)
The output result is:
非零元素索引:
tensor([[0, 1],
[1, 0],
[2, 1]])
Other Extensions