Fix doc and error handling in Argmax

I didn't test the error handling case, but the solution look trivial to me as this is what is done in other cases of error.