Fix exponential gradient reduction without protected attribute (Trust…

…ed-AI#267) Signed-off-by: Drago Plecko <www.plecko@gmail.com>
dplecko · Nov 5, 2021 · 1fd8e71 · 1fd8e71
1 parent 1f1d63a
commit 1fd8e71
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)