Filter properly.

Quantco · Aug 15, 2024 · c0bdcbd · c0bdcbd
1 parent 59554b1
commit c0bdcbd
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/metalearners/xlearner.py b/metalearners/xlearner.py
@@ -566,7 +566,7 @@ def predict_conditional_average_outcomes(
                 ):
                     fold_model = cfe._estimators[fold_index]
                     predict_method = getattr(fold_model, predict_method_name)
-                    fold_estimates = predict_method(X[test_indices])
+                    fold_estimates = predict_method(index_matrix(X, test_indices))
                     conditional_average_outcome_estimates[test_indices] = fold_estimates
 
                 conditional_average_outcomes_list.append(