Merge pull request #301 from credo-ai/fix/DummyClassifier_predict_pro…

…ba_validation Only call check_array if predict or predict proba is provided
credo-ai · Jan 17, 2023 · 45b3069 · 45b3069
2 parents 2327167 + 342319c
commit 45b3069
Showing 1 changed file with 8 additions and 6 deletions.
diff --git a/credoai/artifacts/model/classification_model.py b/credoai/artifacts/model/classification_model.py
@@ -163,12 +163,14 @@ def __init__(
         tags=None,
     ):
         self.model_like = model_like
-        self.predict_output = check_array(
-            predict_output, ensure_2d=False, allow_nd=True
-        )
-        self.predict_proba_output = check_array(
-            predict_proba_output, ensure_2d=False, allow_nd=True
-        )
+        if predict_output is not None:
+            self.predict_output = check_array(
+                predict_output, ensure_2d=False, allow_nd=True
+            )
+        if predict_proba_output is not None:
+            self.predict_proba_output = check_array(
+                predict_proba_output, ensure_2d=False, allow_nd=True
+            )
         self.name = name
         self.tags = tags