From 2a4d6dbcdd4526bdf6068fe9b30d776141b6522a Mon Sep 17 00:00:00 2001 From: Keith Battocchi Date: Mon, 6 Jan 2020 15:15:26 -0500 Subject: [PATCH] Fix linting issue --- econml/drlearner.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/econml/drlearner.py b/econml/drlearner.py index 71ce0c220..504468f0c 100644 --- a/econml/drlearner.py +++ b/econml/drlearner.py @@ -328,7 +328,8 @@ def score(self, Y, T, X=None, W=None, *, nuisances, sample_weight=None, sample_v return np.mean(np.average((Y_pred[..., 1:] - Y_pred[..., [0]] - self.model_cate.predict(X))**2, weights=sample_weight, axis=0)) else: - return np.mean([np.average((Y_pred[..., t] - Y_pred[..., 0] - self.models_cate[t - 1].predict(X))**2, + return np.mean([np.average((Y_pred[..., t] - Y_pred[..., 0] - + self.models_cate[t - 1].predict(X))**2, weights=sample_weight, axis=0) for t in np.arange(1, Y_pred.shape[-1])])