From c0683924efb24f9a6d6a662b1c5572dd6810ca34 Mon Sep 17 00:00:00 2001
From: Sergei Nikolaev <snikolaev@nvidia.com>
Date: Fri, 24 Jun 2016 14:22:42 -0700
Subject: [PATCH] Hot fix preventing dead lock during shutdown

---
 src/caffe/util/gpu_memory.cpp | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/src/caffe/util/gpu_memory.cpp b/src/caffe/util/gpu_memory.cpp
index d83caa20e2e..183fd1ac9dc 100644
--- a/src/caffe/util/gpu_memory.cpp
+++ b/src/caffe/util/gpu_memory.cpp
@@ -37,7 +37,7 @@ void GPUMemory::Manager::init(const vector<int>& gpus, Mode m, bool debug) {
       // Just in case someone installed 'no cleanup' arena before
       delete cub_allocator_;
       cub_allocator_ = new cub::CachingDeviceAllocator(BIN_GROWTH, MIN_BIN,
-          MAX_BIN, MAX_CACHED_BYTES, false, debug_);
+          MAX_BIN, MAX_CACHED_BYTES, true, debug_);
     } catch (...) {
     }
     CHECK_NOTNULL(cub_allocator_);
@@ -117,7 +117,14 @@ void GPUMemory::Manager::deallocate(void* ptr, int device,
   }
   switch (mode_) {
   case CUB_ALLOCATOR:
-    CUDA_CHECK(cub_allocator_->DeviceFree(device, ptr));
+    {
+      int current_device;
+      cudaError_t status = cudaGetDevice(&current_device);
+      // Preventing dead lock while Caffe shutting down.
+      if (status != cudaErrorCudartUnloading) {
+        CUDA_CHECK(cub_allocator_->DeviceFree(device, ptr));
+      }
+    }
     break;
   default:
     CUDA_CHECK(cudaFree(ptr));