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

+2 votes

Best answer

**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)