optimize imports (#61908)
Summary:
Fixes #{issue number}
Pull Request resolved: https://github.com/pytorch/pytorch/pull/61908
Reviewed By: suo
Differential Revision: D29800269
Pulled By: ejguan
fbshipit-source-id: 74ce4414eb6d2a5608df9ec1efdc71e2112aef70
