tf.compat.v1.argmax

Returns the index with the largest value across axes of a tensor. (deprecated arguments)