AutoBatch protect from extreme batch sizes (#9209)

If < 1 or > 1024 set output to default batch size 16. May partially address #9156 Signed-off-by: Glenn Jocher <[email protected]> Signed-off-by: Glenn Jocher <[email protected]>
ultralytics · Aug 29, 2022 · 3c64d89 · 3c64d89
1 parent da22e01
commit 3c64d89
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/utils/autobatch.py b/utils/autobatch.py
@@ -60,8 +60,8 @@ def autobatch(model, imgsz=640, fraction=0.9, batch_size=16):
         i = results.index(None)  # first fail index
         if b >= batch_sizes[i]:  # y intercept above failure point
             b = batch_sizes[max(i - 1, 0)]  # select prior safe point
-    if b < 1:  # zero or negative batch size
-        b = 16
+    if b < 1 or b > 1024:  # b outside of safe range
+        b = batch_size
         LOGGER.warning(f'{prefix}WARNING: ⚠️ CUDA anomaly detected, recommend restart environment and retry command.')
 
     fraction = np.polyval(p, b) / t  # actual fraction predicted