Fix resume_or_load() method in setup_training()

stefanklut · Jan 31, 2024 · cf0b6e1 · cf0b6e1
1 parent 32c6834
commit cf0b6e1
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/main.py b/main.py
@@ -108,10 +108,12 @@ def setup_training(args: argparse.Namespace):
 
         trainer = Trainer(cfg=cfg)
         if not cfg.TRAIN.WEIGHTS:
-            trainer.resume_or_load(resume=cfg.MODEL.RESUME)
+            if cfg.MODEL.RESUME:
+                if not trainer.checkpointer.has_checkpoint():
+                    raise FileNotFoundError(f"No checkpoint found in {cfg.OUTPUT_DIR}")
+                trainer.resume_or_load(resume=cfg.MODEL.RESUME)
         else:
             trainer.checkpointer.load(cfg.TRAIN.WEIGHTS)
-            # if trainer.checkpointer.has_checkpoint():
             trainer.start_iter = trainer.iter + 1
 
         results = trainer.train()