From 63d0c5d6437f09e57e871367286f57adf146b310 Mon Sep 17 00:00:00 2001
From: Hari <harikrishmenon@gmail.com>
Date: Fri, 4 Oct 2024 20:53:47 +0530
Subject: [PATCH] fix: crash when output directory doesn't exist (#364)

Signed-off-by: Harikrishnan Balagopal <harikrishmenon@gmail.com>
---
 tuning/sft_trainer.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tuning/sft_trainer.py b/tuning/sft_trainer.py
index 767fea1b..7b52e3af 100644
--- a/tuning/sft_trainer.py
+++ b/tuning/sft_trainer.py
@@ -637,6 +637,9 @@ def main():
     combined_tracker_configs.file_logger_config = file_logger_config
     combined_tracker_configs.aim_config = aim_config
 
+    if training_args.output_dir:
+        os.makedirs(training_args.output_dir, exist_ok=True)
+        logger.info("using the output directory at %s", training_args.output_dir)
     try:
         trainer, additional_train_info = train(
             model_args=model_args,