"uber#596: Bug fix - Add prediction step to the new case."

AlxClt · Jan 17, 2023 · f4c59be · f4c59be
1 parent bca2b0f
commit f4c59be
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/causalml/inference/meta/xlearner.py b/causalml/inference/meta/xlearner.py
@@ -331,6 +331,9 @@ def estimate_ate(
                 )
             else:
                 p = self._format_p(p, self.t_groups)
+                te, dhat_cs, dhat_ts = self.predict(
+                    X, treatment, y, p=p, return_components=True
+                )
         else:
             te, dhat_cs, dhat_ts = self.fit_predict(
                 X, treatment, y, p, return_components=True