Fix bug which led to crash in kd-test

PiperOrigin-RevId: 701322169
google-research · Nov 29, 2024 · 3e3980b · 3e3980b
1 parent d3d5364
commit 3e3980b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/kauldron/train/trainer_lib.py b/kauldron/train/trainer_lib.py
@@ -399,7 +399,7 @@ def context_specs(self) -> context_lib.Context:
     # Instead just creating the `spec` should be enough.
     m_batch = data_utils.mock_batch_from_elem_spec(elem_spec, elem_sharding)
 
-    context = jax.eval_shape(
+    _, context = jax.eval_shape(
         self.trainstep._step,  # pylint: disable=protected-access
         self.state_specs,
         m_batch,