Internal change

PiperOrigin-RevId: 375539097

Internal change
PiperOrigin-RevId: 375539097
62786702 · Tianqi Liu · A. Unique TensorFlower · c351b6f6 · 62786702
Commit 62786702 authored May 24, 2021 by Tianqi Liu Committed by A. Unique TensorFlower May 24, 2021
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 1 deletion

official/modeling/multitask/train_lib.py official/modeling/multitask/train_lib.py +3 -1

No files found.
--- a/official/modeling/multitask/train_lib.py
+++ b/official/modeling/multitask/train_lib.py
@@ -72,7 +72,9 @@ def run_experiment(*, distribution_strategy: tf.distribute.Strategy,
      evaluator = evaluator_lib.MultiTaskEvaluator(
          task=task,
          model=model,
-          global_step=trainer.global_step if is_training else None)
+          global_step=trainer.global_step if is_training else None,
+          checkpoint_exporter=train_utils.maybe_create_best_ckpt_exporter(
+              params, model_dir))
    else:
      evaluator = None