diff --git a/common/cNet.cpp b/common/cNet.cpp index 86d65b2..5fdea92 100644 --- a/common/cNet.cpp +++ b/common/cNet.cpp @@ -423,6 +423,13 @@ Waifu2x::eWaifu2xError cNet::SetParameter(caffe::NetParameter ¶m, const std: else layer_param->mutable_sigmoid_param()->set_engine(caffe::SigmoidParameter_Engine_CAFFE); } + else if (type == "Pooling") + { + if (process == "cudnn") + layer_param->mutable_pooling_param()->set_engine(caffe::PoolingParameter_Engine_CUDNN); + else + layer_param->mutable_pooling_param()->set_engine(caffe::PoolingParameter_Engine_CAFFE); + } } return Waifu2x::eWaifu2xError_OK;