From b6265da5f38fcc41947d2348ad3a0d09964984bb Mon Sep 17 00:00:00 2001
From: Nikhil Kulkarni <knikhil29@gmail.com>
Date: Thu, 19 May 2022 22:15:35 +0000
Subject: [PATCH] Add more generic OOM

---
 src/sagemaker_server.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/sagemaker_server.cc b/src/sagemaker_server.cc
index 5c71d4c77b..3532ce5173 100644
--- a/src/sagemaker_server.cc
+++ b/src/sagemaker_server.cc
@@ -451,7 +451,7 @@ SagemakerAPIServer::SageMakerMMEHandleLoadError(
       "CUDA out of memory", /* pytorch */
       "CUDA_OUT_OF_MEMORY", /* tensorflow */
       "Out of memory",      /* generic */
-  };
+      "out of memory", "MemoryError"};
 
   EVBufferAddErrorJson(req->buffer_out, load_err);