From cd1c4043531a35c99fd99604ab92bb96a30d242a Mon Sep 17 00:00:00 2001
From: Zeng Jinle <32832641+sneaxiy@users.noreply.github.com>
Date: Thu, 24 Oct 2019 18:53:37 +0800
Subject: [PATCH] refine err msg of allocator, test=develop (#20804)

---
 paddle/fluid/memory/detail/system_allocator.cc | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/paddle/fluid/memory/detail/system_allocator.cc b/paddle/fluid/memory/detail/system_allocator.cc
index 55011179ee..0b14fc6a3f 100644
--- a/paddle/fluid/memory/detail/system_allocator.cc
+++ b/paddle/fluid/memory/detail/system_allocator.cc
@@ -120,7 +120,16 @@ void* GPUAllocator::Alloc(size_t* index, size_t size) {
     gpu_alloc_size_ += size;
     return p;
   } else {
-    PADDLE_ENFORCE_NE(cudaGetLastError(), cudaSuccess);
+    if (result == cudaErrorMemoryAllocation) {
+      result = cudaSuccess;
+    }
+    PADDLE_ENFORCE_CUDA_SUCCESS(result);
+
+    result = cudaGetLastError();
+    if (result == cudaErrorMemoryAllocation) {
+      result = cudaSuccess;
+    }
+    PADDLE_ENFORCE_CUDA_SUCCESS(result);
 
     size_t avail, total;
     platform::GpuMemoryUsage(&avail, &total);
-- 
GitLab