+2 votes
in Programming Languages by (34.9k points)
Which function(s) should I use to find the smallest and largest elements in a given tensor?

1 Answer

+2 votes
by (157k points)
selected by
 
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)


...