Update

[ghstack-poisoned]
pytorch · XuehaiPan · Mar 3, 2025 · Mar 3, 2025 · Mar 3, 2025 · Mar 3, 2025
commit 15f6dfaf95c42c70a557b1889c58ca3e681738fc
@@ -75,7 +75,7 @@
             "Or set the environment variable `PYTORCH_USE_CXX_PYTREE=0`."
         )
 else:
-    cxx = _sys.modules.get("torch.utils._cxx_pytree")
+    cxx = _sys.modules.get("torch.utils._cxx_pytree")  # type: ignore[assignment]
 
 
 _sys.modules[f"{__name__}.python"] = python

diff --git a/torch/utils/pytree/cxx.py b/torch/utils/pytree/cxx.py
@@ -0,0 +1,8 @@
+import sys
+
+import torch.utils._cxx_pytree
+
+
+sys.modules[__name__] = torch.utils._cxx_pytree
+
+del sys, torch
diff --git a/torch/utils/pytree/python.py b/torch/utils/pytree/python.py
@@ -0,0 +1,8 @@
+import sys
+
+import torch.utils._pytree
+
+
+sys.modules[__name__] = torch.utils._pytree
+
+del sys, torch