Merge pull request #20 from achamma723/multi_class

[FIX]: Fix DNN multi_classification case
Parietal-INRIA · Jun 3, 2024 · e7bf244 · e7bf244
2 parents c8282c1 + cd00bfa
commit e7bf244
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/hidimstat/BBI.py b/hidimstat/BBI.py
@@ -764,6 +764,9 @@ def compute_importance(self, X=None, y=None):
         for ind_fold, estimator in enumerate(self.list_estimators):
             if self.type == "DNN":
                 for y_col in range(y[ind_fold].shape[-1]):
+                    _ = self.estimator.encode_outcome(
+                        self.y_train[ind_fold], train=True
+                    )[y_col]
                     y[ind_fold] = self.estimator.encode_outcome(
                         y[ind_fold], train=False
                     )[y_col]