Update fslite/fs/multivariate.py

Co-authored-by: codiumai-pr-agent-pro[bot] <151058649+codiumai-pr-agent-pro[bot]@users.noreply.github.com>
bigbio · Sep 25, 2024 · d5cc974 · d5cc974
1 parent 6ecbaca
commit d5cc974
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/fslite/fs/multivariate.py b/fslite/fs/multivariate.py
@@ -166,7 +166,9 @@ def multivariate_correlation_selector(
         )
 
     # Select feature index to keep
-    selected_features = [i for i in features_indexes if i not in index_to_remove]
+    mask = np.ones(len(features_indexes), dtype=bool)
+    mask[list(index_to_remove)] = False
+    selected_features = np.array(features_indexes)[mask]
 
     return selected_features