diff --git a/main/train.py b/main/train.py
index 3a36e57..65606b5 100644
--- a/main/train.py
+++ b/main/train.py
@@ -86,7 +86,7 @@ def main():
     )
 
     # Set the save directory
-    save_dir = "trained_models_test"
+    save_dir = "trained_models"
     os.makedirs(save_dir, exist_ok=True)
 
     # Load the model from file