Allow NaNs in input X

stanfordmlgroup · Dec 23, 2024 · 77c20fa · 77c20fa
1 parent acf6b27
commit 77c20fa
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/ngboost/ngboost.py b/ngboost/ngboost.py
@@ -1,4 +1,5 @@
 """The NGBoost library"""
+
 # pylint: disable=line-too-long,too-many-instance-attributes,too-many-arguments
 # pylint: disable=unused-argument,too-many-locals,too-many-branches,too-many-statements
 # pylint: disable=unused-variable,invalid-unary-operand-type,attribute-defined-outside-init
@@ -338,7 +339,12 @@ def partial_fit(
             raise ValueError("y cannot be None")
 
         X, Y = check_X_y(
-            X, Y, accept_sparse=True, y_numeric=True, multi_output=self.multi_output
+            X,
+            Y,
+            accept_sparse=True,
+            force_all_finite="allow-nan",
+            multi_output=self.multi_output,
+            y_numeric=True,
         )
 
         self.n_features = X.shape[1]
@@ -353,8 +359,9 @@ def partial_fit(
                 X_val,
                 Y_val,
                 accept_sparse=True,
-                y_numeric=True,
+                force_all_finite="allow-nan",
                 multi_output=self.multi_output,
+                y_numeric=True,
             )
             val_params = self.pred_param(X_val)
             val_loss_list = []