diff --git a/paddle/fluid/memory/detail/system_allocator.cc b/paddle/fluid/memory/detail/system_allocator.cc index 55011179ee8ab824f73aaf79df038660a64a54da..0b14fc6a3fc4a0e4a58b3fac1aafe8ee5c65e8dc 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);