Fix: pytree warning

Xilinx · Dec 31, 2024 · 96700eb · 96700eb
1 parent 39ce837
commit 96700eb
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/brevitas/backport/fx/immutable_collections.py b/src/brevitas/backport/fx/immutable_collections.py
@@ -43,7 +43,7 @@
 
 from typing import Any, Dict, List, Tuple
 
-from torch.utils._pytree import _register_pytree_node
+from torch.utils._pytree import register_pytree_node
 from torch.utils._pytree import Context
 
 from ._compatibility import compatibility
@@ -111,5 +111,5 @@ def _immutable_list_unflatten(values: List[Any], context: Context) -> List[Any]:
     return immutable_list(values)
 
 
-_register_pytree_node(immutable_dict, _immutable_dict_flatten, _immutable_dict_unflatten)
-_register_pytree_node(immutable_list, _immutable_list_flatten, _immutable_list_unflatten)
+register_pytree_node(immutable_dict, _immutable_dict_flatten, _immutable_dict_unflatten)
+register_pytree_node(immutable_list, _immutable_list_flatten, _immutable_list_unflatten)