torch.argmin(): Find the Index of Minimum Value in a Tensor

The torch.argmin() method in PyTorch returns the indices of the minimum value(s) of the flattened tensor or along a dimension. First, it locates the smallest value from the tensor and then returns its index, not the value.

The above figure shows how the .argmin() method works on a 1D tensor.

The type of the output tensor of indices is torch.int64.

If multiple minimum values exist, it returns the index of the first occurrence.

The torch.argmin() method is different from the torch.min() method because it returns just the indices of the minimum value, while the .min() method returns the minimum value and its indices.

Syntax

torch.argmin(input, dim=None, keepdim=False)

Parameters

Argument	Description
input (Tensor)	It is the input tensor from which we will find the indices of the minimum value.
dim (int, optional)	If None, it flattens the tensor. If it is provided, it is the dimension to reduce.
keepdim (bool, optional)	If True, retains the reduced dimension with size 1. Default: False.

1D Tensor

Let’s define a 1D tensor and find the index of minima. The index starts with 0 in a tensor.

import torch

tensor = torch.tensor([3, 4, 1, 2, 5])

min_idx = torch.argmin(tensor)

print(min_idx)

# Output: tensor(2)

In the above code, the minimum value is 1 and its index is 2 because it is the 3rd element.

Multi-Dimensional Tensor (Flattened)

In the above figure, we found the index of a global minimum value from the 2D tensor.

Let’s implement the figure into a Python program.

import torch

tensor_2d = torch.tensor([[21, 4, 10],
                          [2, 5, 19]])

min_idx = torch.argmin(tensor_2d)

print(min_idx)

# Output: tensor(3)

If you don’t pass the “dim” argument in a multidimensional tensor, the argmin() method will first flatten the tensor to 1D and then return the index of the minimum value.

In our case, the minimum value is at index 3, which is value 2. So, it returns a tensor(3) as an output.

Along a Dimension (2D Tensor)

For a 2D tensor, if you pass dim=0, it will find the index of the minimum value along the columns. It operates down the rows, column-wise.

If you pass dim=1, it will find the minimum value index along the tensor rows. It operates down the columns, row-wise.

import torch

tensor_2d = torch.tensor([[21, 4, 10],
                          [2, 5, 19]])

min_along_columns = torch.argmin(tensor_2d, dim=0)

print(min_along_columns)
# Output: tensor([1, 0, 0])

min_along_rows = torch.argmin(tensor_2d, dim=1)

print(min_along_rows)
# Output: tensor([1, 0])

If we check column-wise, the first column’s two values are 21 and 2, and 2 is the minimum, and its index is 1. Here, make sure to count the index from 0 column-wise.

The second column’s two values are 4 and 5. The 4 is minimum, and its index is 0.

The third column’s three values are 10 and 19. The 19 is the minimum, and its index is 0.

Hence, column-wise, we got the output tensor([1, 0, 0]).

The same logic applies to row-wise.

The first row has three values, 21, 4, and 10. The minimum value is 10, and its index is 1.

The second row has three values, 2, 5, and 19. The minimum value is 2, and its index is 0.

Hence, row-wise, we got the output tensor([1, 0]). And there are only two rows.

Using keepdim

The keepdim=True argument ensures that the result keeps the reduced dimension as size 1. So, if the input tensor has a [2, 2] shape, the output tensor remains 2D with shape (2, 1) instead of being flattened to 1D (2, ).

import torch

tensor_2d = torch.tensor([[21, 4, 10],
                          [2, 5, 19]])

reduced_indices = torch.argmin(tensor_2d, dim=1, keepdim=True)

print(reduced_indices)

# Output:
# tensor([[1],
#         [0]])

We can evaluate the output like this:

1. Row 0: [21, 4, 10] → min value is 4 at index 1
2. Row 1: [2, 5, 19] → min value is 2 at index 0

Empty Tensor

If you pass the empty tensor to the torch.argmin() method, it will throw: IndexError: argmin(): Expected reduction dim to be specified for input.numel() == 0 error and our program is crashed.

import torch

empty_tensor = torch.tensor([])

indices = torch.argmin(empty_tensor)

print(indices)

# Output:
# IndexError: argmin(): Expected reduction dim to be specified for input.numel() == 0.

We can handle this type of error using the “try/except” block to prevent the program from crashing.

import torch

empty_tensor = torch.tensor([])

try:
    idx = torch.argmin(empty_tensor)
except IndexError as e:
    print(e)

# Output: argmin(): Expected reduction dim to be specified for input.numel() == 0.