[Resubmit] Fix argmin/max bug (#39576)
Summary:
Fix https://github.com/pytorch/pytorch/issues/38922
See previous PR: https://github.com/pytorch/pytorch/pull/38946
cc: ngimel
Pull Request resolved: https://github.com/pytorch/pytorch/pull/39576
Differential Revision: D21906490
Pulled By: ngimel
fbshipit-source-id: f3bfb4e14c4cee60a1e3b80c049945ce85f9f494
