[Refactor] Upgrade pytree import (#573)

pytorch · Nov 24, 2023 · df1ee89 · df1ee89
1 parent 42bf143
commit df1ee89
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tensordict/_pytree.py b/tensordict/_pytree.py
@@ -11,7 +11,7 @@
     TensorDict,
 )
 
-from torch.utils._pytree import _register_pytree_node, Context
+from torch.utils._pytree import Context, register_pytree_node
 
 PYTREE_REGISTERED_TDS = (
     LazyStackedTensorDict,
@@ -99,7 +99,7 @@ def _tensordictdict_unflatten(values: List[Any], context: Context) -> Dict[Any,
 
 
 for cls in PYTREE_REGISTERED_TDS:
-    _register_pytree_node(
+    register_pytree_node(
         cls,
         _tensordict_flatten,
         _tensordictdict_unflatten,