fix TypeError in feature pruning

CoDS-GCS · Jul 6, 2023 · cd08be6 · cd08be6
1 parent 4d8a3e6
commit cd08be6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/interface/apis.py b/interface/apis.py
@@ -280,7 +280,7 @@ def compute_feature_importance(train: pd.DataFrame, test: pd.DataFrame, theta1:
 
         def compute_feature_correlation(train: pd.DataFrame, test: pd.DataFrame, theta2: float):
             nonlocal pruning_info
-            features = train[set(numerical_features).intersection(set(pruning_info.keys()))]
+            features = train[list(set(numerical_features).intersection(set(pruning_info.keys())))]
             if target in list(features.columns):
                 features = train.drop(target, axis=1)
             features_to_discard = set()