DOC: fix import error (#38921)
Summary:
Fixes errors when importing the module. The import is used by sphinx in documentation builds.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/38921
Differential Revision: D21722144
Pulled By: ezyang
fbshipit-source-id: 5f31d4750325f1753de93754a009006cbc13655e
