undo pytree.py change

Lightning-AI · Jun 24, 2024 · 57e1249 · 57e1249
1 parent b45dbd0
commit 57e1249
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/thunder/core/pytree.py b/thunder/core/pytree.py
@@ -49,4 +49,4 @@ def tree_unflatten(values, spec):
     return optree.tree_unflatten(spec, values)
 
 
-__all__ = ["tree_flatten", "tree_unflatten", "tree_map", "tree_flatten_with_dataclass"]
+__all__ = ["tree_flatten", "tree_unflatten", "tree_map"]