fix(api): handle more out-of-memory errors in the workers

ssube · Mar 16, 2023 · c8c5e9f · c8c5e9f
1 parent dfba5f0
commit c8c5e9f
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/api/onnx_web/worker/worker.py b/api/onnx_web/worker/worker.py
@@ -60,7 +60,8 @@ def worker_main(context: WorkerContext, server: ServerContext):
             )
             exit(EXIT_ERROR)
         except Exception as e:
-            if "Failed to allocate memory" in str(e) or "CUDA out of memory" in str(e):
+            e_str = str(e)
+            if "Failed to allocate memory" in e_str or "out of memory" in e_str:
                 logger.error("detected out-of-memory error, exiting: %s", e)
                 exit(EXIT_MEMORY)
             else: