Tensor.argmax()

Returns the indices of the maximum values along an axis.

Usage

from tinygrad.tensor import Tensor
 
tensor = Tensor([1, 2, 3, 4, 5, 6])
 
print(tensor.argmax())

5.0

;