diff --git a/learned_optimization/tree_utils.py b/learned_optimization/tree_utils.py
index 5467c60..d346f02 100644
--- a/learned_optimization/tree_utils.py
+++ b/learned_optimization/tree_utils.py
@@ -243,6 +243,6 @@ def partition_unflatten(unflattener: PartitionUnflatten,
   to_fill = [None for _ in keys]
   for name, part in zip(names, part_values):
     for n, p in zip(name, part):
-      to_fill[unmap[n]] = p
+      to_fill[unmap[n]] = p  # pyrefly: ignore[unsupported-operation]
 
   return jax.tree_util.tree_unflatten(struct, to_fill)