fix train_test_prediction_drift reduce fails on value of 0 (#2346)

deepchecks · Feb 21, 2023 · ad707d2 · ad707d2
1 parent 8121b7d
commit ad707d2
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/deepchecks/tabular/checks/model_evaluation/train_test_prediction_drift.py b/deepchecks/tabular/checks/model_evaluation/train_test_prediction_drift.py
@@ -11,6 +11,7 @@
 """Module contains Train Test label Drift check."""
 
 import typing as t
+from numbers import Number
 
 import numpy as np
 import pandas as pd
@@ -247,7 +248,7 @@ def run_logic(self, context: Context) -> CheckResult:
 
     def reduce_output(self, check_result: CheckResult) -> t.Dict[str, float]:
         """Return prediction drift score."""
-        if isinstance(check_result.value['Drift score'], float):
+        if isinstance(check_result.value['Drift score'], Number):
             return {'Prediction Drift Score': check_result.value['Drift score']}
 
         drift_values = list(check_result.value['Drift score'].values())