From 2dfb3e7a952399105b2f57d1a4b67d59dccc1936 Mon Sep 17 00:00:00 2001 From: Sergei Nikolaev Date: Tue, 24 May 2016 17:54:06 -0700 Subject: [PATCH] Fix for OOM when CUB succeeds to free cache and reallocate --- src/caffe/util/gpu_memory.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/caffe/util/gpu_memory.cpp b/src/caffe/util/gpu_memory.cpp index 8cf79c94d2f..790af689926 100644 --- 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));