fix

openvinotoolkit · Nov 30, 2023 · ad39bd0 · ad39bd0
1 parent 00b0939
commit ad39bd0
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/model_api/cpp/models/src/segmentation_model.cpp b/model_api/cpp/models/src/segmentation_model.cpp
@@ -253,7 +253,8 @@ void SegmentationModel::prepareInputsOutputs(std::shared_ptr<ov::Model>& model)
 
 std::unique_ptr<ResultBase> SegmentationModel::postprocess(InferenceResult& infResult) {
     const auto& inputImgSize = infResult.internalModelData->asRef<InternalImageModelData>();
-    const auto& outTensor = infResult.outputsData[outputNames[0]];
+    const auto& outputName = outputNames[0] == feature_vector_name ? outputNames[1] : outputNames[0];
+    const auto& outTensor = infResult.outputsData[outputName];
     const ov::Shape& outputShape = outTensor.get_shape();
     const ov::Layout& outputLayout = getLayoutFromShape(outputShape);
     size_t outChannels = ov::layout::has_channels(outputLayout) ?