clean up

allenai · Dec 7, 2023 · a44217a · a44217a
1 parent 569381a
commit a44217a
Showing 1 changed file with 4 additions and 9 deletions.
diff --git a/olmo/train.py b/olmo/train.py
@@ -910,15 +910,10 @@ def on_trace_ready(p):
                 log.info("Training loop complete")
 
         # Save final unsharded model-only checkpoint.
-        if not canceled:
-            if self.cfg.save_interval_unsharded is not None:
-                log.info("Saving final unsharded checkpoint...")
-                checkpoint_path, _ = self.save_checkpoint(CheckpointType.unsharded)
-                log.info(f"Unsharded checkpoint saved to {checkpoint_path}")
-            else:
-                log.info("Saving final sharded checkpoint...")
-                checkpoint_path, _ = self.save_checkpoint(CheckpointType.sharded)
-                log.info(f"Checkpoint saved to {checkpoint_path}")
+        if not canceled and self.cfg.save_interval_unsharded is not None:
+            log.info("Saving final unsharded model checkpoint...")
+            checkpoint_path, _ = self.save_checkpoint(CheckpointType.unsharded)
+            log.info(f"Unsharded checkpoint saved to {checkpoint_path}")
 
     def close(self, exit_code: int = 0) -> None:
         if self.indices_file is not None: