[fix] tmp fix for import issue in dtensor (#119582)
a temporary fix for S394053 which is likely caused by backward incompatible `import` introduced in D53437243. It's yet to be understood why this may cause an issue but let's forward "fix" it first then draft a follow up diff for a right fix.
Differential Revision: [D53621345](https://our.internmc.facebook.com/intern/diff/D53621345/)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/119582
Approved by: https://github.com/tianyu-l