From 137cbc9184123140d1e98ee7abb88fdc9c42197a Mon Sep 17 00:00:00 2001 From: Ahmed Ahmed Date: Thu, 15 Feb 2024 16:04:28 -0800 Subject: [PATCH] fix --- src/levanter/main/train_lm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/levanter/main/train_lm.py b/src/levanter/main/train_lm.py index 67157523b..34db8849f 100644 --- 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.")