tf.compat.v1.arg_max

Returns the index with the lar