fix partitioning

pfackeldey · May 8, 2024 · c6f6361 · c6f6361
1 parent b36faaf
commit c6f6361
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/evermore/parameter.py b/src/evermore/parameter.py
@@ -146,7 +146,7 @@ def _replace_value(leaf: Any) -> Any:
             leaf = eqx.tree_at(lambda p: p.value, leaf, not leaf.frozen)
         return leaf
 
-    return params_map(_replace_value, filter_spec)
+    return jtu.tree_map(_replace_value, filter_spec, is_leaf=is_parameter)
 
 
 def partition(tree: PyTree) -> tuple[PyTree, PyTree]: