check that checkpoint folder exists before training model

openspyrit · Jul 4, 2024 · aa17397 · aa17397
1 parent 6628dee
commit aa17397
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/spyrit/core/train.py b/spyrit/core/train.py
@@ -178,6 +178,10 @@ def train_model(
     train_info["train"] = []
     train_info["val"] = []
 
+    # check that the folder `root` exists if do_checkpoint > 0
+    if (do_checkpoint > 0) and (not os.path.exists(root)):
+        raise ValueError(f"Folder {root} not found")
+
     # Set tensorboard writer
     if tb_path:
         samples, _ = next(iter(dataloaders["val"]))