From b2eb406197d208c135d51a2347a5b073ec3490f7 Mon Sep 17 00:00:00 2001 From: Sean Sube Date: Wed, 15 Mar 2023 08:51:29 -0500 Subject: [PATCH] fix(api): handle CUDA memory errors in workers --- api/onnx_web/worker/worker.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/api/onnx_web/worker/worker.py b/api/onnx_web/worker/worker.py index 3afefe65..34dd65e4 100644 --- a/api/onnx_web/worker/worker.py +++ b/api/onnx_web/worker/worker.py @@ -60,7 +60,7 @@ def worker_main(context: WorkerContext, server: ServerContext): ) exit(EXIT_ERROR) except Exception as e: - if "Failed to allocate memory" in str(e): + if "Failed to allocate memory" in str(e) or "CUDA out of memory" in str(e): logger.error("detected out-of-memory error, exiting: %s", e) exit(EXIT_MEMORY) else: