Partial revert of #38144 to fix ROCm CI. (#38363)
Summary:
CC ezyang xw285cornell
Pull Request resolved: https://github.com/pytorch/pytorch/pull/38363
Differential Revision: D21539778
Pulled By: ezyang
fbshipit-source-id: 0f7d3b8e3b30ab4d5992f1c13aa8d48069796a8d
