fix test

jerryming1995 · Jul 7, 2023 · fb633d6 · fb633d6
1 parent c98cf42
commit fb633d6
Showing 1 changed file with 9 additions and 5 deletions.
diff --git a/pycaret/internal/pycaret_experiment/supervised_experiment.py b/pycaret/internal/pycaret_experiment/supervised_experiment.py
@@ -4933,7 +4933,8 @@ def encode_labels(label_encoder, labels: pd.Series) -> pd.Series:
             else:
                 data = self._set_index(self._prepare_dataset(data))
                 target = None
-
+            X_test_untransformed = data
+            y_test_untransformed = target
             data = data[X_columns]  # Ignore all columns but the originals
             if preprocess:
                 X_test_ = pipeline.transform(
@@ -4951,11 +4952,14 @@ def encode_labels(label_encoder, labels: pd.Series) -> pd.Series:
                 X_test_ = data
                 y_test_ = target
 
-            X_test_untransformed = data[data.index.isin(X_test_.index)]
+            # Align number of rows with output of transformation
+            X_test_untransformed = X_test_untransformed[
+                X_test_untransformed.index.isin(X_test_.index)
+            ]
             if target is not None:
-                y_test_untransformed = target[target.index.isin(y_test_.index)]
-            else:
-                y_test_untransformed = target
+                y_test_untransformed = y_test_untransformed[
+                    y_test_untransformed.index.isin(X_test_.index)
+                ]
 
         # prediction starts here
         if isinstance(estimator, CustomProbabilityThresholdClassifier):