diff --git a/train.py b/train.py
index 794cadd..4830601 100644
--- a/train.py
+++ b/train.py
@@ -36,8 +36,13 @@ def train_model():
         cache='disk',        # Changed to disk caching for deterministic results
     )
     
+     # Create model directory if it doesn't exist
+    import os
+    os.makedirs('model/weights', exist_ok=True)
+
     # Save the trained model
     model.save('model/weights/best.pt')
 
+
 if __name__ == '__main__':
     train_model()
\ No newline at end of file