Fix docstring for argmax (#23775)
Summary:
Fixes https://github.com/pytorch/pytorch/issues/23757
Pull Request resolved: https://github.com/pytorch/pytorch/pull/23775
Differential Revision: D16644692
Pulled By: soumith
fbshipit-source-id: d759bb85f73383021e4657325dbac79913042ad2
