Internal change

PiperOrigin-RevId: 377381104

Internal change
PiperOrigin-RevId: 377381104
88a6cf61 · Le Hou · A. Unique TensorFlower · 069bdd28 · 88a6cf61
Commit 88a6cf61 authored Jun 03, 2021 by Le Hou Committed by A. Unique TensorFlower Jun 03, 2021
Show whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

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

No files found.
--- a/official/modeling/progressive/trainer.py
+++ b/official/modeling/progressive/trainer.py
@@ -284,8 +284,11 @@ class ProgressiveTrainer(trainer_lib.Trainer):
          checkpoint_interval=checkpoint_interval,
      )

+    # Make sure we export the last checkpoint.
+    last_checkpoint = (
+        self.global_step.numpy() == self._config.trainer.train_steps)
    checkpoint_path = self._export_ckpt_manager.save(
        checkpoint_number=self.global_step.numpy(),
-        check_interval=True)
+        check_interval=not last_checkpoint)
    if checkpoint_path:
      logging.info('Checkpoints exported: %s.', checkpoint_path)