autocreate save_ckpt_dir folder if not exists

2023-01-12 12:25:09 -05:00 · 2023-01-12 12:25:09 -05:00 · ff86644e8c
parent 0222a984a2
commit ff86644e8c
1 changed files with 3 additions and 0 deletions
--- 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):