add keep_dims to the reshape condition

andrei-cv · Jun 7, 2021 · 37d022e · 37d022e
1 parent dd47b89
commit 37d022e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/inference-engine/src/cldnn_engine/ops/reduce.cpp b/inference-engine/src/cldnn_engine/ops/reduce.cpp
@@ -81,7 +81,7 @@ void CreateReduceOp(Program& p, const std::shared_ptr<ngraph::Node>& op, cldnn::
 
     auto resultLayerName = layerName;
     auto out_dims = op->get_output_shape(0).size();
-    if (out_dims == 3 && rank >= 4) {
+    if (out_dims == 3 && !keep_dims && rank >= 4) {
         resultLayerName = layerName + "_reshape";
         auto out_shape = op->get_output_shape(0);
         cldnn::tensor outTensor;