argmax for half datatype fix (#28787)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/28787
Stack from [ghstack](https://github.com/ezyang/ghstack):
* **#28787 argmax for half datatype fix**
Test Plan: Imported from OSS
Differential Revision: D18194420
Pulled By: pbelevich
fbshipit-source-id: d2abec1ea8a9ce3a93aec5a2c5bba57d163197e6
