IVGCVSW-2834 Fix Input TensorInfo in Quantization DataSet

Signed-off-by: Narumol Prangnawarat <narumol.prangnawarat@arm.com>
Change-Id: I694124aa15ef2fd368d5e6e36da8b942ae270850

diff --git a/src/armnnQuantizer/QuantizationDataSet.cpp b/src/armnnQuantizer/QuantizationDataSet.cpp
index d225883..ac819dc 100644 (file)
--- a/src/armnnQuantizer/QuantizationDataSet.cpp
+++ b/src/armnnQuantizer/QuantizationDataSet.cpp
@@ -51,7 +51,7 @@ void InputLayerVisitor::VisitInputLayer(const armnn::IConnectableLayer* layer,
                                         armnn::LayerBindingId id,
                                         const char* name)
 {
-    m_TensorInfos.emplace(id, layer->GetInputSlot(0).GetConnection()->GetTensorInfo());
+    m_TensorInfos.emplace(id, layer->GetOutputSlot(0).GetTensorInfo());
 }
 
 armnn::TensorInfo InputLayerVisitor::GetTensorInfo(armnn::LayerBindingId layerBindingId)