Features to drop bug fix (Trusted-AI#60)

fix bug in handling of features_to_drop
shanky3011 · Dec 11, 2018 · f5485c3 · f5485c3
1 parent 2a5eff0
commit f5485c3
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/aif360/datasets/standard_dataset.py b/aif360/datasets/standard_dataset.py
@@ -95,7 +95,8 @@ def __init__(self, df, label_name, favorable_classes,
         if instance_weights_name:
             keep |= set([instance_weights_name])
         df = df[sorted(keep - set(features_to_drop), key=df.columns.get_loc)]
-
+        categorical_features = sorted(set(categorical_features) - set(features_to_drop), key=df.columns.get_loc)
+
         # 4. Remove any rows that have missing data.
         dropped = df.dropna()
         count = df.shape[0] - dropped.shape[0]