From 05ba95e10630bd294c568c5ac63046638109bd47 Mon Sep 17 00:00:00 2001
From: Chen Weihang <chenweihang@baidu.com>
Date: Tue, 7 Feb 2023 20:14:09 +0800
Subject: [PATCH] Fix bad alloc exception throw error (#50266)

* fix bad alloc exp error

* remove needless files
---
 paddle/fluid/pybind/exception.cc | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/paddle/fluid/pybind/exception.cc b/paddle/fluid/pybind/exception.cc
index 934a9ef97f..2fac903c8d 100644
--- a/paddle/fluid/pybind/exception.cc
+++ b/paddle/fluid/pybind/exception.cc
@@ -14,6 +14,7 @@ limitations under the License. */
 
 #include "paddle/fluid/pybind/exception.h"
 
+#include "paddle/fluid/memory/allocation/allocator.h"
 #include "paddle/phi/api/ext/exception.h"
 namespace paddle {
 namespace pybind {
@@ -41,6 +42,8 @@ void BindException(pybind11::module* m) {
       if (p) std::rethrow_exception(p);
     } catch (const platform::EOFException& e) {
       eof(e.what());
+    } catch (const memory::allocation::BadAlloc& e) {
+      PyErr_SetString(PyExc_MemoryError, e.what());
     } catch (const platform::EnforceNotMet& e) {
       switch (e.code()) {
         case paddle::platform::error::INVALID_ARGUMENT:
@@ -91,6 +94,8 @@ void ThrowExceptionToPython(std::exception_ptr p) {
     if (p) std::rethrow_exception(p);
   } catch (const platform::EOFException& e) {
     PyErr_SetString(EOFExceptionException, e.what());
+  } catch (const memory::allocation::BadAlloc& e) {
+    PyErr_SetString(PyExc_MemoryError, e.what());
   } catch (const platform::EnforceNotMet& e) {
     switch (e.code()) {
       case paddle::platform::error::INVALID_ARGUMENT:
-- 
GitLab