[PyTorch] How to find the index of minimum and maximum elements in a tensor

Question

Which function(s) should I use to find the smallest and largest elements in a given tensor?

Answer 1

argmax() and argmin() functions can be used to find the maximum and minimum elements in a given tensor, respectively. These functions are similar to Numpy's argmax() and argmin().

Here is an example:

>>> import torch
>>> a=torch.randn(10)
>>> a
tensor([ 1.0062,  1.0328,  0.1764, -0.5594, -0.0861,  0.4747,  0.1676,  1.7200,
        -2.5906, -0.6808])
>>> torch.argmax(a)
tensor(7)
>>> torch.argmin(a)
tensor(8)