Internal change

PiperOrigin-RevId: 352912327

Internal change
PiperOrigin-RevId: 352912327
e0f818da · Le Hou · A. Unique TensorFlower · 488d20bc · e0f818da
Commit e0f818da authored Jan 20, 2021 by Le Hou Committed by A. Unique TensorFlower Jan 20, 2021
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

official/modeling/progressive/trainer.py official/modeling/progressive/trainer.py +2 -1

No files found.
--- a/official/modeling/progressive/trainer.py
+++ b/official/modeling/progressive/trainer.py
@@ -267,7 +267,8 @@ class ProgressiveTrainer(trainer_lib.Trainer):
      step_interval = self.config.trainer.checkpoint_interval
    else:
      step_interval = self.config.trainer.export_checkpoint_interval
-    if global_step_np % step_interval != 0:
+    if global_step_np % step_interval != 0 and (
+        global_step_np < self._config.trainer.train_steps):
      logging.info('Not exporting checkpoints in global step: %d.',
                   global_step_np)
      return