diff options
Diffstat (limited to 'src/armnn/OutputHandler.hpp')
-rw-r--r-- | src/armnn/OutputHandler.hpp | 5 |
1 files changed, 5 insertions, 0 deletions
diff --git a/src/armnn/OutputHandler.hpp b/src/armnn/OutputHandler.hpp index 41a49af031..3fd2519ed5 100644 --- a/src/armnn/OutputHandler.hpp +++ b/src/armnn/OutputHandler.hpp @@ -50,10 +50,15 @@ public: void SetData(std::unique_ptr<ITensorHandle> data) { m_TensorHandle = std::move(data); } + void SetAllocatedData() { m_AllocatedTensorHandle = std::move(m_TensorHandle); } + + void UseAllocatedData() { m_TensorHandle = std::move(m_AllocatedTensorHandle); } + /// @brief Returns true if SetTensorInfo() has been called at least once on this. bool IsTensorInfoSet() const { return m_bTensorInfoSet; } private: std::unique_ptr<ITensorHandle> m_TensorHandle; + std::unique_ptr<ITensorHandle> m_AllocatedTensorHandle; TensorInfo m_TensorInfo; bool m_bTensorInfoSet = false; }; |