FIX add condition on impute_missing

neurodata · Jul 9, 2024 · 3e2cda8 · 3e2cda8
1 parent 9126a9d
commit 3e2cda8
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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: