CUDAが使えない環境でCPUでUpResNet10を動作させるとクラッシュするのを修正 resolved #126

lltcggie · Oct 25, 2018 · 7d1973c · 7d1973c
1 parent de6de69
commit 7d1973c
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/common/cNet.cpp b/common/cNet.cpp
@@ -390,6 +390,13 @@ Waifu2x::eWaifu2xError cNet::SetParameter(caffe::NetParameter &param, 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;