refine err msg of allocator, test=develop (#20804)

6 years ago · cd1c404353
parent 370f0345b6
commit cd1c404353
1 changed files with 10 additions and 1 deletions
--- 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);