Fix exponential gradient reduction without protected attribute (Trust…

…ed-AI#267)
Illia-Kryvoviaz · Sep 6, 2021 · 68b40b3 · 68b40b3
1 parent 966db55
commit 68b40b3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/aif360/sklearn/inprocessing/exponentiated_gradient_reduction.py b/aif360/sklearn/inprocessing/exponentiated_gradient_reduction.py
@@ -144,6 +144,6 @@ 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)
 
         return self.model._pmf_predict(X)