Remained changes of #43578 (#43921)
Summary:
Not full https://github.com/pytorch/pytorch/issues/43578 was merged. This PR is the remained part.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/43921
Reviewed By: ailzhang
Differential Revision: D23438504
Pulled By: mruberry
fbshipit-source-id: 9c5e26346dfc423b7a440b8a986420a27349090f
