From aa1739794e206e9441ccc228a0b77a2be6e087f9 Mon Sep 17 00:00:00 2001
From: romainphan <romainphan.rp@gmail.com>
Date: Thu, 4 Jul 2024 15:41:15 +0200
Subject: [PATCH] check that checkpoint folder exists before training model

---
 spyrit/core/train.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/spyrit/core/train.py b/spyrit/core/train.py
index 6c95e6ee..f7be269e 100644
--- 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"]))