Fix addOutput (#1888)

diff --git a/inference-engine/src/inference_engine/cnn_network_ngraph_impl.cpp b/inference-engine/src/inference_engine/cnn_network_ngraph_impl.cpp
index 8137e1b..2efd5cf 100644 (file)
--- a/inference-engine/src/inference_engine/cnn_network_ngraph_impl.cpp
+++ b/inference-engine/src/inference_engine/cnn_network_ngraph_impl.cpp
@@ -260,6 +260,8 @@ StatusCode CNNNetworkNGraphImpl::addOutput(const std::string& layerName, size_t
 void CNNNetworkNGraphImpl::addOutput(const ::ngraph::Output<::ngraph::Node> & output) {
     auto dataName = ngraph::op::util::create_ie_output_name(output);
     DataPtr data;
+    if (_data.count(dataName))
+        data = _data[dataName];
     createDataForResult(output, dataName, data);
     _data[dataName] = data;
     _outputData[dataName] = data;

diff --git a/inference-engine/tests/functional/inference_engine/cnn_network/cnn_ngraph_impl_tests.cpp b/inference-engine/tests/functional/inference_engine/cnn_network/cnn_ngraph_impl_tests.cpp
index 02e138b..11c68af 100644 (file)
--- a/inference-engine/tests/functional/inference_engine/cnn_network/cnn_ngraph_impl_tests.cpp
+++ b/inference-engine/tests/functional/inference_engine/cnn_network/cnn_ngraph_impl_tests.cpp
@@ -753,4 +753,37 @@ TEST(CNNNGraphImplTests, addOutput) {
     ASSERT_EQ(outputs["reshape"]->getLayout(), InferenceEngine::Layout::NCHW);
 }
 
+TEST(CNNNGraphImplTests, addOutputForParameter) {
+    std::shared_ptr<ngraph::Function> ngraph;
+    {
+        ngraph::PartialShape shape({1, 3, 22, 22});
+        ngraph::element::Type type(ngraph::element::Type_t::f32);
+        auto param = std::make_shared<ngraph::opset3::Parameter>(type, shape);
+        param->set_friendly_name("param");
+        auto relu = std::make_shared<ngraph::opset3::Relu>(param);
+        auto result = std::make_shared<ngraph::op::Result>(relu);
+
+        ngraph::ParameterVector params = {param};
+        ngraph::ResultVector results = {result};
+
+        ngraph = std::make_shared<ngraph::Function>(results, params);
+    }
+
+    CNNNetwork cnnNetwork(ngraph);
+    cnnNetwork.addOutput("param");
+    {
+        auto output_info = cnnNetwork.getOutputsInfo();
+        ASSERT_EQ(output_info.count("param"), 1);
+        ASSERT_EQ(output_info["param"]->getTensorDesc().getDims(), SizeVector({1, 3, 22, 22}));
+    }
+
+    cnnNetwork.reshape({{"param", SizeVector({1, 3, 32, 64})}});
+    cnnNetwork.addOutput("param");
+    {
+        auto output_info = cnnNetwork.getOutputsInfo();
+        ASSERT_EQ(output_info.count("param"), 1);
+        ASSERT_EQ(output_info["param"]->getTensorDesc().getDims(), SizeVector({1, 3, 32, 64}));
+    }
+}
+
 IE_SUPPRESS_DEPRECATED_END