From 96141092f1bcccaea9a5cd1cdef68fa9d6c49163 Mon Sep 17 00:00:00 2001
From: Adeyeri Michael <67843464+Mikehade@users.noreply.github.com>
Date: Wed, 22 Mar 2023 10:26:58 +0100
Subject: [PATCH] Update stratify parameter

---
 whitebox/analytics/models/pipelines.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/whitebox/analytics/models/pipelines.py b/whitebox/analytics/models/pipelines.py
index 02d63ba..ceea0ab 100644
--- a/whitebox/analytics/models/pipelines.py
+++ b/whitebox/analytics/models/pipelines.py
@@ -95,7 +95,7 @@ def create_multiclass_classification_training_model_pipeline(
     """
 
     X_train, X_test, y_train, y_test = train_test_split(
-        X, Y, test_size=0.3, random_state=0
+        X, Y, test_size=0.3, random_state=0, stratify = Y
     )
     d_train = lgb.Dataset(X_train, label=y_train)
     """