fix: allow more runtime errors in find_batch_size()

This attempts to fix #358.
kreshuklab · Nov 18, 2024 · d8ffe1d · d8ffe1d
1 parent 36e9445
commit d8ffe1d
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/plantseg/functionals/prediction/utils/size_finder.py b/plantseg/functionals/prediction/utils/size_finder.py
@@ -170,6 +170,10 @@ def find_batch_size(
                 if "out of memory" in str(e):
                     batch_size //= 2
                     break
+                elif "less than INT_MAX elements" in str(e):
+                    logger.warning(f"Encountered '{e}', unexpected but continuing with reduced batch size.")
+                    batch_size //= 2
+                    break
                 else:
                     raise
             finally: