Internal change

PiperOrigin-RevId: 350687522

official/core/base_trainer.py

@@ -161,7 +161,8 @@ class Trainer(orbit.StandardTrainer, orbit.StandardEvaluator):
           self,
           eval_dataset,
           options=orbit.StandardEvaluatorOptions(
-              use_tf_function=config.trainer.eval_tf_function))
+              use_tf_function=config.trainer.eval_tf_function,
+              use_tf_while_loop=config.trainer.eval_tf_while_loop))
 
   def _validate_params(self, config):
     r"""Validates if the configuration object passed to the Trainer.

official/core/config_definitions.py

@@ -192,6 +192,7 @@ class TrainerConfig(base_config.Config):
   train_tf_while_loop: bool = True
   train_tf_function: bool = True
   eval_tf_function: bool = True
+  eval_tf_while_loop: bool = False
   allow_tpu_summary: bool = False
   # Trainer intervals.
   steps_per_loop: int = 1000