Fix typo in argmax pseudocode

Axis is a dimension index of the input tensor, but is used to index the
output tensor.

Signed-off-by: Thomas Preud'homme <thomas.preudhomme@arm.com>
Change-Id: Ie50f3921153e01437fe763b5d2b6cfaa06049f25