Update train.py

train.py

@@ -5,6 +5,7 @@ from sklearn.tree import DecisionTreeClassifier
 from sklearn.model_selection import GridSearchCV
 import joblib
 import wandb
+import os
 
 
 
@@ -41,6 +42,8 @@ def run_training(max_depth=10, min_samples_split=2):
 
     # Save the best model
     best_model = grid_search.best_estimator_
+    # Create the 'model' directory if it doesn't exist
+    os.makedirs('model', exist_ok=True)
     joblib.dump(best_model, "model/best_model.pkl")
 
     print("Training completed and best model saved.")