From 3e2cda8191648dff05e17f9fa64465c74dcb6b2a Mon Sep 17 00:00:00 2001 From: Haoyin Xu Date: Tue, 9 Jul 2024 16:22:42 -0400 Subject: [PATCH] FIX add condition on impute_missing --- sktree/ensemble/_honest_forest.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/sktree/ensemble/_honest_forest.py b/sktree/ensemble/_honest_forest.py index 3d662c9a0..f9ba3ec5d 100644 --- a/sktree/ensemble/_honest_forest.py +++ b/sktree/ensemble/_honest_forest.py @@ -672,7 +672,8 @@ def _predict_proba(self, X, indices=None, impute_missing=None): zero_mask = posteriors.sum(2) == 0 posteriors[~zero_mask] /= posteriors[~zero_mask].sum(1, keepdims=True) - posteriors[zero_mask] = impute_missing + if impute_missing is not None: + posteriors[zero_mask] = impute_missing # preserve shape of multi-outputs if self.n_outputs_ > 1: