Adding data seen to the reported statistics on the evaluator in the e…

…xamples. PiperOrigin-RevId: 508676532
google-deepmind · Feb 10, 2023 · d6f14ad · d6f14ad
1 parent c4e5765
commit d6f14ad
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/examples/training.py b/examples/training.py
@@ -435,7 +435,7 @@ def _evaluate_single_batch(
   ) -> Dict[str, chex.Array]:
     """Evaluates a single batch."""
 
-    del global_step, opt_state  # This might be used in subclasses
+    del global_step  # This might be used in subclasses
 
     func_args = kfac_jax.optimizer.make_func_args(
         params=params,
@@ -447,6 +447,7 @@ def _evaluate_single_batch(
     )
     loss, stats = self.eval_model_func(*func_args)
     stats["loss"] = loss
+    stats["data_seen"] = opt_state.data_seen
 
     return kfac_jax.utils.pmean_if_pmap(stats, "eval_axis")