Log a message when an iterator is exhausted during a Orbit training/evaluation.

PiperOrigin-RevId: 413973623

orbit/utils/loop_fns.py

@@ -14,6 +14,7 @@
 
 """Utilities for creating loop functions."""
 
+from absl import logging
 from orbit.utils import tpu_summaries
 
 import tensorflow as tf
@@ -65,8 +66,8 @@ def create_loop_fn(step_fn):
       The final state returned by `reduce_fn`, or `None` if `state` and
       `reduce_fn` are not provided.
     """
-    try:
     step = 0
+    try:
       # To make sure the OutOfRangeError exception can be handled well under
       # async remote eager, we need to wrap the loop body in `async_scope`.
       with tf.experimental.async_scope():
@@ -77,6 +78,7 @@ def create_loop_fn(step_fn):
           step += 1
         return state
     except (StopIteration, tf.errors.OutOfRangeError):
+      logging.info("The dataset iterator is exhausted after %d steps.", step)
       tf.experimental.async_clear_error()
       return state