pytree

Genentech · Aug 19, 2024 · 3c188e2 · 3c188e2
1 parent 60f37cc
commit 3c188e2
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/src/beignet/func/__dataclass.py b/src/beignet/func/__dataclass.py
@@ -1,7 +1,7 @@
 import dataclasses
 from typing import List, Tuple, Type, TypeVar, Iterable
 
-import optree
+from torch.utils._pytree import register_pytree_node
 
 T = TypeVar("T")
 
@@ -42,11 +42,10 @@ def _iterable_to_cls(meta, data) -> T:
 
         return dataclass_cls(**kwargs)
 
-    optree.register_pytree_node(
+    register_pytree_node(
         dataclass_cls,
         _iterate_cls,
         _iterable_to_cls,
-        namespace="beignet.func",
     )
 
     return dataclass_cls