pytorch / 2 / generated / torch.argmin.html

torch.argmin

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

Returns the indices of the minimum value(s) of the flattened tensor or along a dimension

This is the second value returned by torch.min(). See its documentation for the exact semantics of this method.

Note

If there are multiple minimal values then the indices of the first minimal value are returned.

Parameters
  • input (Tensor) – the input tensor.
  • dim (int) – the dimension to reduce. If None, the argmin of the flattened input is returned.
  • keepdim (bool) – whether the output tensor has dim retained or not..

Example:

>>> a = torch.randn(4, 4)
>>> a
tensor([[ 0.1139,  0.2254, -0.1381,  0.3687],
        [ 1.0100, -1.1975, -0.0102, -0.4732],
        [-0.9240,  0.1207, -0.7506, -1.0213],
        [ 1.7809, -1.2960,  0.9384,  0.1438]])
>>> torch.argmin(a)
tensor(13)
>>> torch.argmin(a, dim=1)
tensor([ 2,  1,  3,  1])
>>> torch.argmin(a, dim=1, keepdim=True)
tensor([[2],
        [1],
        [3],
        [1]])

© 2024, PyTorch Contributors
PyTorch has a BSD-style license, as found in the LICENSE file.
https://pytorch.org/docs/2.1/generated/torch.argmin.html