cuDNNのdllがない環境でGPUモードで実行すると強制終了するバグを修正

This commit is contained in:
lltcggie 2016-07-06 21:30:02 +09:00
parent 7f15cd6fb4
commit 2ea99797c6

View File

@ -362,6 +362,13 @@ Waifu2x::eWaifu2xError cNet::SetParameter(caffe::NetParameter &param, const std:
else else
layer_param->mutable_convolution_param()->set_engine(caffe::ConvolutionParameter_Engine_CAFFE); layer_param->mutable_convolution_param()->set_engine(caffe::ConvolutionParameter_Engine_CAFFE);
} }
else if (type == "Deconvolution")
{
if (process == "cudnn")
layer_param->mutable_convolution_param()->set_engine(caffe::ConvolutionParameter_Engine_CUDNN);
else
layer_param->mutable_convolution_param()->set_engine(caffe::ConvolutionParameter_Engine_CAFFE);
}
else if (type == "ReLU") else if (type == "ReLU")
{ {
if (process == "cudnn") if (process == "cudnn")