From d4abb2bec28b0a852ad7c86c2d1bfdc856414f34 Mon Sep 17 00:00:00 2001
From: "Malte S. Kurz" <malte.s.kurz@gmail.com>
Date: Fri, 4 Jun 2021 11:01:09 +0200
Subject: [PATCH] by initializing with nan instead of zero, we no longer get
 divide by zero warnings when running IRM with ATTE score without
 cross-fitting

---
 doubleml/double_ml_irm.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/doubleml/double_ml_irm.py b/doubleml/double_ml_irm.py
index da918ef1..98c0d4e4 100644
--- a/doubleml/double_ml_irm.py
+++ b/doubleml/double_ml_irm.py
@@ -192,7 +192,7 @@ def _score_elements(self, y, d, g_hat0, g_hat1, m_hat, smpls):
         # fraction of treated for ATTE
         p_hat = None
         if self.score == 'ATTE':
-            p_hat = np.zeros_like(d, dtype='float64')
+            p_hat = np.full_like(d, np.nan, dtype='float64')
             for _, test_index in smpls:
                 p_hat[test_index] = np.mean(d[test_index])