Merge pull request #22 from ag-tcm/patch-1

Update BorutaShap.py
Ekeany · Sep 23, 2020 · 9aee98b · 9aee98b
2 parents 320c28b + e443dea
commit 9aee98b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/BorutaShap.py b/src/BorutaShap.py
@@ -204,7 +204,7 @@ def Check_if_chose_train_or_test_and_train_model(self):
 
         if self.train_or_test.lower() == 'test':
             # keeping the same naming convenetion as to not add complexit later on
-            self.X_boruta_train, self.X_boruta, self.y_train, self.y_test = train_test_split(self.X_boruta,
+            self.X_boruta_train, self.X_boruta_test, self.y_train, self.y_test = train_test_split(self.X_boruta,
                                                                                             self.y,
                                                                                             test_size=0.3,
                                                                                             random_state=self.random_state)
@@ -1014,4 +1014,4 @@ def load_data(data_type='classification'):
         raise ValueError("No data_type was specified, use either 'classification' or 'regression'")
 
 
-    return X, y
+    return X, y