From 569aa498d05a558727b7cf27648b6874f70d54c2 Mon Sep 17 00:00:00 2001
From: Sergei Nikolaev <snikolaev@nvidia.com>
Date: Mon, 18 Jan 2016 15:27:40 -0800
Subject: [PATCH] cuDNN and pooling fixes

---
 src/caffe/layer_factory.cpp                | 7 +++++++
 src/caffe/layers/cudnn_batch_norm_layer.cu | 4 ++++
 2 files changed, 11 insertions(+)

diff --git a/src/caffe/layer_factory.cpp b/src/caffe/layer_factory.cpp
index 0d1625602db..2237175daa8 100644
--- a/src/caffe/layer_factory.cpp
+++ b/src/caffe/layer_factory.cpp
@@ -78,7 +78,14 @@ shared_ptr<Layer<Dtype> > GetPoolingLayer(const LayerParameter& param) {
 #ifdef USE_CUDNN
   } else if (engine == PoolingParameter_Engine_CUDNN) {
     PoolingParameter p_param = param.pooling_param();
+
+// FIXME We assume that precision issue will be fixed in cuDNN v5.0
+#if CUDNN_VERSION >= 5000
     return shared_ptr<Layer<Dtype> >(new CuDNNPoolingLayer<Dtype>(param));
+#else
+    return shared_ptr<Layer<Dtype> >(new PoolingLayer<Dtype>(param));
+#endif
+
 #endif
   } else {
     LOG(FATAL) << "Layer " << param.name() << " has unknown engine.";
diff --git a/src/caffe/layers/cudnn_batch_norm_layer.cu b/src/caffe/layers/cudnn_batch_norm_layer.cu
index bc1ee25154d..bacb80e9b09 100644
--- a/src/caffe/layers/cudnn_batch_norm_layer.cu
+++ b/src/caffe/layers/cudnn_batch_norm_layer.cu
@@ -85,6 +85,10 @@ void CuDNNBatchNormLayer<Dtype>::Backward_gpu(
       mode_,
       cudnn::dataType<Dtype>::one,
       cudnn::dataType<Dtype>::zero,
+#if CUDNN_VERSION >= 4005
+      cudnn::dataType<Dtype>::one,
+      cudnn::dataType<Dtype>::zero,
+#endif
       bottom_desc_,
       bottom_data,
       bottom_desc_,