Fix D23995953 import.
Summary: https://github.com/pytorch/pytorch/pull/45511 could not be properly imported
Test Plan: See https://github.com/pytorch/pytorch/pull/45511
Reviewed By: zhangguanheng66
Differential Revision: D23995953
fbshipit-source-id: a6224a67d54617ddf34c2392e65f2142c4e78ea4