diff --git a/paddle/fluid/pybind/exception.cc b/paddle/fluid/pybind/exception.cc
index 934a9ef97fb15486d5dfa7cfe4d6c0c89b076d5c..2fac903c8dff35a188d8466a5035346b56c46039 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: