diff options
Diffstat (limited to 'src/armnn/QuantizerVisitor.hpp')
| -rw-r--r-- | src/armnn/QuantizerVisitor.hpp | 5 |
1 files changed, 4 insertions, 1 deletions
diff --git a/src/armnn/QuantizerVisitor.hpp b/src/armnn/QuantizerVisitor.hpp index 2682663047..1751229435 100644 --- a/src/armnn/QuantizerVisitor.hpp +++ b/src/armnn/QuantizerVisitor.hpp @@ -7,6 +7,7 @@ #include "armnn/LayerVisitorBase.hpp" #include "StaticRangeVisitor.hpp" +#include "NetworkQuantizationScheme.hpp" #include <armnn/INetwork.hpp> #include <armnn/Types.hpp> @@ -24,7 +25,7 @@ class StaticRangeVisitor; class QuantizerVisitor : public LayerVisitorBase<VisitorNoThrowPolicy> { public: - QuantizerVisitor(const RangeTracker& rangeTracker); + QuantizerVisitor(const RangeTracker& rangeTracker, const IQuantizationScheme* quantizationScheme); ~QuantizerVisitor() = default; /// Functions to quantize the individual layers, overridden from ILayerVisitor @@ -129,6 +130,8 @@ private: /// Mapping from guid to layer in quantized network std::unordered_map<LayerGuid, IConnectableLayer*> m_QuantizedGuidToLayerMap; + + const IQuantizationScheme* m_QuantizationScheme; }; } //namespace armnn |