fix

stanford-crfm · Feb 16, 2024 · 137cbc9 · 137cbc9
1 parent 1bacc9b
commit 137cbc9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/levanter/main/train_lm.py b/src/levanter/main/train_lm.py
@@ -150,7 +150,7 @@ def add_floats(x, y):
 
                 # what is the f here?
                 logger.info(f"Interpolating between the two models with alpha={alpha}")
-                merged_model = named_jit(jax.tree_util.tree_map, add_floats)(model, model_2)
+                merged_model = named_jit(lambda m1, m2: jax.tree_util.tree_map(add_floats, m1, m2), donate=True)(model, model_2)
                 state = dataclasses.replace(state, model=model)
             else:
                 logger.info("No checkpoint found. Starting from scratch.")