diff --git a/common/cNet.cpp b/common/cNet.cpp index 0e84f3c..902858d 100644 --- a/common/cNet.cpp +++ b/common/cNet.cpp @@ -390,6 +390,13 @@ Waifu2x::eWaifu2xError cNet::SetParameter(caffe::NetParameter ¶m, const std: else layer_param->mutable_relu_param()->set_engine(caffe::ReLUParameter_Engine_CAFFE); } + else if (type == "Sigmoid") + { + if (process == "cudnn") + layer_param->mutable_sigmoid_param()->set_engine(caffe::SigmoidParameter_Engine_CUDNN); + else + layer_param->mutable_sigmoid_param()->set_engine(caffe::SigmoidParameter_Engine_CAFFE); + } } return Waifu2x::eWaifu2xError_OK;