Revert "Partial revert of #38144 to fix ROCm CI. (#38363)" (#38380)
Summary:
The changes in this file broke ROCm and got reverted in https://github.com/pytorch/pytorch/issues/38363. This PR brings it back with ROCm fixed.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/38380
Differential Revision: D21549632
Pulled By: ezyang
fbshipit-source-id: 68498aba70e651352d58fd0c865e71420dbf900a
