Add missing line to build_tree_from_paths

Author: samanklesaria

flax/jax_utils.py

@@ -352,6 +352,6 @@ def build_tree_from_paths(paths_and_leaves: list[tuple[jax.tree_util.KeyPath, An
           current = current[k]
 
       # Set the leaf value
-      final_key = path[-1]
-      current[_path_ix(final_key)] = leaf
+      current.is_list = isinstance(path[-1], jax.tree_util.SequenceKey)
+      current[_path_ix(path[-1])] = leaf
   return _to_pytree(root)