correcting typo bug causing WMT validation metric failure

PiperOrigin-RevId: 711475895
google · Jan 2, 2025 · 559b15d · 559b15d
1 parent 0788186
commit 559b15d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/init2winit/mt_eval/mt_callback.py b/init2winit/mt_eval/mt_callback.py
@@ -88,7 +88,7 @@ def __init__(self,
     self.callback_config = merged_callback_config
 
     self._validate_callback_config()
-    self.evaluate_batch_pmapped = jax.jit(
+    self.evaluate_batch_jitted = jax.jit(
         model.evaluate_batch, donate_argnums=(2,)
     )
     self.batch_stats = batch_stats