used the checkpoint path as the temporary path for huggingface trainer.

TL-System · Nov 9, 2023 · b57ebc7 · b57ebc7
1 parent d3e79b9
commit b57ebc7
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/examples/split_learning/llm_split_learning/split_learning_trainer.py b/examples/split_learning/llm_split_learning/split_learning_trainer.py
@@ -104,7 +104,10 @@ def __init__(self, model=None, callbacks=None):
         parser = HfArgumentParser(TrainingArguments)
 
         (self.training_args,) = parser.parse_args_into_dataclasses(
-            args=["--report_to=none"]
+            args=[
+                "--output_dir=" + Config.params["checkpoint_path"],
+                "--report_to=none",
+            ]
         )
 
     # Redesign the evaluation stage.