From ff86644e8cf89007d9a3b576cc873275686e4b79 Mon Sep 17 00:00:00 2001
From: Victor Hall <victor.charles.hall@gmail.com>
Date: Thu, 12 Jan 2023 12:25:09 -0500
Subject: [PATCH] autocreate save_ckpt_dir folder if not exists

---
 train.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/train.py b/train.py
index 9f47336..b6fe28a 100644
--- a/train.py
+++ b/train.py
@@ -267,6 +267,9 @@ def setup_args(args):
         args.lr = args.lr * (total_batch_size**0.55)
         logging.info(f"{Fore.CYAN} * Scaling learning rate {tmp_lr} by {total_batch_size**0.5}, new value: {args.lr}{Style.RESET_ALL}")
 
+    if not os.path.exists(args.save_ckpt_dir):
+        os.makedirs(args.save_ckpt_dir)
+
     return args
 
 def main(args):