diff --git a/src/caffe/layer_factory.cpp b/src/caffe/layer_factory.cpp index 4d912d28351..e967bd6181c 100644 --- a/src/caffe/layer_factory.cpp +++ b/src/caffe/layer_factory.cpp @@ -91,7 +91,16 @@ shared_ptr > GetPoolingLayer(const LayerParameter& param) { << "Using Caffe's own pooling layer."; return shared_ptr >(new PoolingLayer(param)); } - return shared_ptr >(new CuDNNPoolingLayer(param)); + // CuDNN assumes layers are not being modified in place, thus + // breaking our index tracking for updates in some cases in Caffe. + // Until there is a workaround in Caffe (index management) or + // cuDNN, use Caffe layer to max pooling, or don't use in place + // layers after max pooling layers + if (param.pooling_param().pool() == PoolingParameter_PoolMethod_MAX) { + return shared_ptr >(new PoolingLayer(param)); + } else { + return shared_ptr >(new CuDNNPoolingLayer(param)); + } #endif } else { LOG(FATAL) << "Layer " << param.name() << " has unknown engine.";