From b454c5c65efb238c130b042ace390b2bc7f0bf75 Mon Sep 17 00:00:00 2001
From: Finn Williams <Finn.Williams@arm.com>
Date: Tue, 9 Feb 2021 15:56:23 +0000
Subject: IVGCVSW-4893 Refactor ILayerVisitor using unified interface strategy.

Signed-off-by: Jan Eilers <jan.eilers@arm.com>
Signed-off-by: Finn Williams <Finn.Williams@arm.com>
Signed-off-by: Francis Murtagh <francis.murtagh@arm.com>
Change-Id: Id7bc8255a8e3f9e5aac65d510bec8a559bf37246
---
 src/armnnQuantizer/QuantizationDataSet.hpp | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

(limited to 'src/armnnQuantizer/QuantizationDataSet.hpp')

diff --git a/src/armnnQuantizer/QuantizationDataSet.hpp b/src/armnnQuantizer/QuantizationDataSet.hpp
index 3a97630ccf..47b893a7f7 100644
--- a/src/armnnQuantizer/QuantizationDataSet.hpp
+++ b/src/armnnQuantizer/QuantizationDataSet.hpp
@@ -42,6 +42,22 @@ private:
     std::string m_CsvFilePath;
 };
 
+/// Visitor class implementation to gather the TensorInfo for LayerBindingID for creation of ConstTensor for Refine.
+class InputLayerStrategy : public armnn::IStrategy
+{
+public:
+    virtual void ExecuteStrategy(const armnn::IConnectableLayer* layer,
+                                 const armnn::BaseDescriptor& descriptor,
+                                 const std::vector<armnn::ConstTensor>& constants,
+                                 const char* name,
+                                 const armnn::LayerBindingId id = 0) override;
+
+    armnn::TensorInfo GetTensorInfo(armnn::LayerBindingId);
+private:
+    std::map<armnn::LayerBindingId, armnn::TensorInfo> m_TensorInfos;
+};
+
+
 /// Visitor class implementation to gather the TensorInfo for LayerBindingID for creation of ConstTensor for Refine.
 class InputLayerVisitor : public armnn::LayerVisitorBase<armnn::VisitorNoThrowPolicy>
 {
-- 
cgit v1.2.1