Fix out-of-sample evaluate.

Quantco · Aug 15, 2024 · fe16b75 · fe16b75
1 parent c0bdcbd
commit fe16b75
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/metalearners/xlearner.py b/metalearners/xlearner.py
@@ -374,8 +374,8 @@ def evaluate(
             tv_imputed_te_control, tv_imputed_te_treatment = self._pseudo_outcome(
                 y, w, treatment_variant, conditional_average_outcome_estimates
             )
-            imputed_te_control.append(tv_imputed_te_control)
-            imputed_te_treatment.append(tv_imputed_te_treatment)
+            imputed_te_control.append(tv_imputed_te_control[w == 0])
+            imputed_te_treatment.append(tv_imputed_te_treatment[w == treatment_variant])
 
         te_treatment_evaluation = _evaluate_model_kind(
             self._treatment_models[TREATMENT_EFFECT_MODEL],