fix multiclass test

MAIF · May 6, 2024 · 1b1789a · 1b1789a
1 parent c026fd5
commit 1b1789a
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/tests/integration_tests/test_contributions_multiclass.py b/tests/integration_tests/test_contributions_multiclass.py
@@ -97,7 +97,10 @@ def test_rank_contributions_1(self):
         model.fit(self.x_train, self.y_train)
         explainer = shap.TreeExplainer(model)
         shap_values = explainer.shap_values(self.x_test)
-        slist = [pd.DataFrame(data=tab, index=self.x_test.index, columns=self.x_test.columns) for tab in shap_values]
+        slist = [
+            pd.DataFrame(data=shap_values[:, :, i], index=self.x_test.index, columns=self.x_test.columns)
+            for i in range(3)
+        ]
 
         for i in range(3):
             s_ord, x_ord, s_dict = rank_contributions(slist[i], pd.DataFrame(data=self.x_test))