Getting Started
Python API
argmax()
Programming Guide
Returns the maximum index of all elements in the input tensor along the provided axis
input
axis
input – the input values
axis – the dimension along which the reduction should be done