diff --git a/training/scripts/prepare_dataset.py b/training/scripts/prepare_dataset.py
index e7313aa..b10230e 100755
--- a/training/scripts/prepare_dataset.py
+++ b/training/scripts/prepare_dataset.py
@@ -62,11 +62,11 @@ def prepare_dataset(input_file, output_file, test_size=0.05):
     test_data = formatted[split_point:]
     
     # Save
-    with open(output_file.parent / "train.jsonl", "w") as f:
+    with open(output_dir / "train.jsonl", "w") as f:
         for item in train_data:
             f.write(json.dumps(item) + "\n")
     
-    with open(output_file.parent / "test.jsonl", "w") as f:
+    with open(output_dir / "test.jsonl", "w") as f:
         for item in test_data:
             f.write(json.dumps(item) + "\n")