diff --git a/src/caffe/layer_factory.cpp b/src/caffe/layer_factory.cpp index f14253a510e..5f6d46a727c 100644 --- a/src/caffe/layer_factory.cpp +++ b/src/caffe/layer_factory.cpp @@ -50,7 +50,7 @@ shared_ptr > GetConvolutionLayer( if (engine == ConvolutionParameter_Engine_DEFAULT) { engine = ConvolutionParameter_Engine_CAFFE; #ifdef USE_CUDNN - if (!use_dilation) { + if (!use_dilation && Caffe::mode() == Caffe::GPU) { engine = ConvolutionParameter_Engine_CUDNN; } #endif @@ -80,7 +80,9 @@ shared_ptr > GetPoolingLayer(const LayerParameter& param) { if (engine == PoolingParameter_Engine_DEFAULT) { engine = PoolingParameter_Engine_CAFFE; #ifdef USE_CUDNN - engine = PoolingParameter_Engine_CUDNN; + if (Caffe::mode() == Caffe::GPU) { + engine = PoolingParameter_Engine_CUDNN; + } #endif } if (engine == PoolingParameter_Engine_CAFFE) { @@ -118,7 +120,9 @@ shared_ptr > GetLRNLayer(const LayerParameter& param) { if (engine == LRNParameter_Engine_DEFAULT) { #ifdef USE_CUDNN - engine = LRNParameter_Engine_CUDNN; + if (Caffe::mode() == Caffe::GPU) { + engine = LRNParameter_Engine_CUDNN; + } #else engine = LRNParameter_Engine_CAFFE; #endif @@ -156,7 +160,9 @@ shared_ptr > GetReLULayer(const LayerParameter& param) { if (engine == ReLUParameter_Engine_DEFAULT) { engine = ReLUParameter_Engine_CAFFE; #ifdef USE_CUDNN - engine = ReLUParameter_Engine_CUDNN; + if (Caffe::mode() == Caffe::GPU) { + engine = ReLUParameter_Engine_CUDNN; + } #endif } if (engine == ReLUParameter_Engine_CAFFE) { @@ -180,7 +186,9 @@ shared_ptr > GetSigmoidLayer(const LayerParameter& param) { if (engine == SigmoidParameter_Engine_DEFAULT) { engine = SigmoidParameter_Engine_CAFFE; #ifdef USE_CUDNN - engine = SigmoidParameter_Engine_CUDNN; + if (Caffe::mode() == Caffe::GPU) { + engine = SigmoidParameter_Engine_CUDNN; + } #endif } if (engine == SigmoidParameter_Engine_CAFFE) { @@ -204,7 +212,9 @@ shared_ptr > GetSoftmaxLayer(const LayerParameter& param) { if (engine == SoftmaxParameter_Engine_DEFAULT) { engine = SoftmaxParameter_Engine_CAFFE; #ifdef USE_CUDNN - engine = SoftmaxParameter_Engine_CUDNN; + if (Caffe::mode() == Caffe::GPU) { + engine = SoftmaxParameter_Engine_CUDNN; + } #endif } if (engine == SoftmaxParameter_Engine_CAFFE) { @@ -228,7 +238,9 @@ shared_ptr > GetTanHLayer(const LayerParameter& param) { if (engine == TanHParameter_Engine_DEFAULT) { engine = TanHParameter_Engine_CAFFE; #ifdef USE_CUDNN - engine = TanHParameter_Engine_CUDNN; + if (Caffe::mode() == Caffe::GPU) { + engine = TanHParameter_Engine_CUDNN; + } #endif } if (engine == TanHParameter_Engine_CAFFE) {