pytorch

313557a6 - Add missing import (#72840)

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)