Update train.py

swarajpande5 · Apr 10, 2024 · 67ee633 · 67ee633
1 parent bf95ed8
commit 67ee633
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/train.py b/train.py
@@ -15,7 +15,7 @@
 X_train, X_test, y_train, y_test = train_test_split(df, y, test_size=0.2, random_state=seed)
 
 # TRAINING
-clf = RandomForestClassifier(max_depth=5, random_state=seed)
+clf = RandomForestClassifier(max_depth=10, random_state=seed)
 clf.fit(X_train, y_train)
 y_pred = clf.predict(X_test)
 
@@ -26,4 +26,4 @@
 cm = confusion_matrix(y_test, y_pred, labels=clf.classes_)
 disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=clf.classes_)
 disp.plot()
-plt.savefig("confmat.png", dpi=120)
+plt.savefig("confmat.png", dpi=120)