diff --git a/training/scripts/prepare_dataset.py b/training/scripts/prepare_dataset.py
index c81fc4d..e4e24eb 100755
--- a/training/scripts/prepare_dataset.py
+++ b/training/scripts/prepare_dataset.py
@@ -66,17 +66,17 @@ def prepare_dataset(input_file, output_file, test_size=0.05):
     test_data = formatted[split_point:]
     
     # Save
-    with open(output_dir / "train.jsonl", "w") as f:
+    with open(output_file / "train.jsonl", "w") as f:
         for item in train_data:
             f.write(json.dumps(item) + "\n")
     
-    with open(output_dir / "test.jsonl", "w") as f:
+    with open(output_file / "test.jsonl", "w") as f:
         for item in test_data:
             f.write(json.dumps(item) + "\n")
     
     print(f"Train: {len(train_data)} examples")
     print(f"Test: {len(test_data)} examples")
-    print(f"Saved to {output_file.parent}")
+    print(f"Saved to {output_file}")
 
 
 def main():