From 0785dcd6f7fc5700935955eba2ac3892fb8bd86e Mon Sep 17 00:00:00 2001 From: kklein Date: Thu, 15 Aug 2024 21:26:23 +0200 Subject: [PATCH] Fix in-sample evaluate. --- metalearners/xlearner.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/metalearners/xlearner.py b/metalearners/xlearner.py index fd76719..e15c842 100644 --- 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],