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

1 Answer

+2 votes
by (45.8k points)

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)


...