From 255137b6e8ca68bf317afd426068cc5beeb4726e Mon Sep 17 00:00:00 2001
From: Noah Harasz <ndharasz@live.com>
Date: Tue, 12 Sep 2023 16:11:46 -0700
Subject: [PATCH] check that predictions are between 0 and 1

---
 predict.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/predict.py b/predict.py
index d695394..cbd70d0 100644
--- a/predict.py
+++ b/predict.py
@@ -171,6 +171,11 @@ def main(args):
             elif predictions.isna().any().any():
                 logging.error("Pickle function returned at least 1 NaN prediction")
                 exit_with_help(1)
+            elif not (predictions.between(0, 1).all().all()):
+                logging.error(
+                    "Pickle function returned invalid predictions. Ensure values are between 0 and 1."
+                )
+                exit_with_help(1)
         except TypeError as e:
             logging.error(f"Pickle function is invalid - {e}")
             if args.debug: