add mu_dtype guard

stanford-crfm · May 1, 2024 · b942d60 · b942d60
1 parent ae6c90c
commit b942d60
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/levanter/optim/schedulefree_adam.py b/src/levanter/optim/schedulefree_adam.py
@@ -86,7 +86,7 @@ def _adam_gradient_transform(
       A `GradientTransformation` object.
     """
 
-    mu_dtype = jax.canonicalize_dtype(mu_dtype)
+    mu_dtype = jax.canonicalize_dtype(mu_dtype) if mu_dtype is not None else None
 
     def init_fn(params):
         z = jax.tree_util.tree_map(jnp.copy, params)  # schedule-free z