Fix predict_proba with drop_prot_attr

close #363
Trusted-AI · Sep 3, 2022 · ed452b0 · ed452b0
1 parent 297379a
commit ed452b0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/aif360/sklearn/inprocessing/grid_search_reduction.py b/aif360/sklearn/inprocessing/grid_search_reduction.py
@@ -193,7 +193,7 @@ def predict_proba(self, X):
             ``self.classes_``.
         """
         if self.drop_prot_attr:
-            X = X.drop(self.prot_attr)
+            X = X.drop(self.prot_attr, axis=1)
 
         if isinstance(self.model_.constraints, red.ClassificationMoment):
             return self.model_.predict_proba(X)