Add missing import (#72840)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/72840
Reviewed By: H-Huang
Differential Revision: D34242612
Pulled By: albanD
fbshipit-source-id: 3dd34de96dbf1ae8f3c3ea45888d211d95862c49
(cherry picked from commit d2650ffa75dbba315daeb0e7cdf0fcb56f3584e1)
