diff --git a/include/caffe/util/cudnn.hpp b/include/caffe/util/cudnn.hpp index b531dd5fa7a..8a7e17c6cd4 100644 --- a/include/caffe/util/cudnn.hpp +++ b/include/caffe/util/cudnn.hpp @@ -7,6 +7,9 @@ #include "caffe/common.hpp" #include "caffe/proto/caffe.pb.h" +#define CUDNN_VERSION_MIN(major, minor, patch) \ + (CUDNN_VERSION >= (major * 1000 + minor * 100 + patch)) + #define CUDNN_CHECK(condition) \ do { \ cudnnStatus_t status = condition; \ diff --git a/src/caffe/layers/cudnn_conv_layer.cu b/src/caffe/layers/cudnn_conv_layer.cu index 1990e932a70..42c4fd0260c 100644 --- a/src/caffe/layers/cudnn_conv_layer.cu +++ b/src/caffe/layers/cudnn_conv_layer.cu @@ -30,11 +30,19 @@ void CuDNNConvolutionLayer::Forward_gpu( // Bias. if (this->bias_term_) { const Dtype* bias_data = this->blobs_[1]->gpu_data(); +#if CUDNN_VERSION_MIN(4, 0, 0) + CUDNN_CHECK(cudnnAddTensor(handle_[g], + cudnn::dataType::one, + bias_desc_, bias_data + bias_offset_ * g, + cudnn::dataType::one, + top_descs_[i], top_data + top_offset_ * g)); +#else CUDNN_CHECK(cudnnAddTensor(handle_[g], CUDNN_ADD_SAME_C, cudnn::dataType::one, bias_desc_, bias_data + bias_offset_ * g, cudnn::dataType::one, top_descs_[i], top_data + top_offset_ * g)); +#endif } }