diff --git a/sepes/_src/tree_mask.py b/sepes/_src/tree_mask.py
index 482975f..5f14370 100644
--- a/sepes/_src/tree_mask.py
+++ b/sepes/_src/tree_mask.py
@@ -202,9 +202,6 @@ def _tree_mask_map(
     *,
     is_leaf: Callable[[Any], None] | None = None,
 ):
-    treelib = sepes._src.backend.treelib
-    # apply func to leaves satisfying mask pytree/condtion
-    _, lhsdef = treelib.flatten(tree, is_leaf=is_leaf)
 
     if not isinstance(cond, Callable):
         # a callable that accepts a leaf and returns a boolean
@@ -214,6 +211,8 @@ def _tree_mask_map(
             f" Got {cond=} and {tree=}."
         )
 
+    treelib = sepes._src.backend.treelib
+
     def map_func(x):
         return func(x) if cond(x) else x