PyTorch torch.argwhere() method “returns a tensor containing the indices of all non-zero elements of the input.” Each row in the result contains the indices of a non-zero element in input. The result is sorted lexicographically, with the last index changing the fastest (C-style).
Syntax
torch.argwhere(input)Parameters
input (Tensor): The input tensor.
Example 1: Basic usage with a 2D tensor
import torch
tensor_1 = torch.tensor([[0, 1, 0], [2, 0, 3], [0, 0, 0]])
indices_1 = torch.argwhere(tensor_1)
print("Original Tensor:")
print(tensor_1)
print("\nIndices of Non-Zero Elements:")
print(indices_1)
Output
Example 2: Using the method with a 3D tensor
import torch
tensor_2 = torch.tensor([
[[0, 1], [2, 0]],
[[3, 0], [0, 4]]
])
indices_2 = torch.argwhere(tensor_2)
print("Original Tensor:")
print(tensor_2)
print("\nIndices of Non-Zero Elements:")
print(indices_2)Output
Example 3: Using the method with a 1D tensor
import torch
tensor_3 = torch.tensor([0, 1, 2, 0, 3, 0])
indices_3 = torch.argwhere(tensor_3)
print("Original Tensor:")
print(tensor_3)
print("\nIndices of Non-Zero Elements:")
print(indices_3)Output
Original Tensor:
tensor([0, 1, 2, 0, 3, 0])
Indices of Non-Zero Elements:
tensor([[1],
[2],
[4]])For a 1D tensor, the returned indices are simply the positions of the non-zero elements in the tensor.
Related posts
Krunal Lathiya is a seasoned Computer Science expert with over eight years in the tech industry. He boasts deep knowledge in Data Science and Machine Learning. Versed in Python, JavaScript, PHP, R, and Golang. Skilled in frameworks like Angular and React and platforms such as Node.js. His expertise spans both front-end and back-end development. His proficiency in the Machine Learning frameworks like PyTorch and Tensorflow is a testament to his versatility and commitment to the craft.