Resolve comments in #44354. (#45150)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/45150
Test Plan: Imported from OSS
Reviewed By: bhosmer
Differential Revision: D23846796
Pulled By: ailzhang
fbshipit-source-id: 7bef89d833848ac3f8993c4c037acf1d4f2ca674
