Fix for OOM when CUB succeeds to free cache and reallocate

NVIDIA · May 25, 2016 · 2dfb3e7 · 2dfb3e7
1 parent 19492ba
commit 2dfb3e7
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/caffe/util/gpu_memory.cpp b/src/caffe/util/gpu_memory.cpp
@@ -82,9 +82,9 @@ void GPUMemoryManager::allocate(void** ptr, size_t size, cudaStream_t stream) {
       }
       // Retry once
       CUDA_CHECK(cub_allocator->DeviceAllocate(ptr, size, stream));
-      // If retry succeeds we need to clean up last error
-      cudaGetLastError();
     }
+    // If retry succeeds we need to clean up last error
+    cudaGetLastError();
     break;
   default:
     CUDA_CHECK(cudaMalloc(ptr, size));