fix: crash when output directory doesn't exist (foundation-model-stac…

…k#364) Signed-off-by: Harikrishnan Balagopal <[email protected]>
red-hat-data-services · Oct 4, 2024 · 63d0c5d · 63d0c5d
1 parent 7b97e9e
commit 63d0c5d
Showing 1 changed file with 3 additions and 0 deletions.
diff --git 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,