diff options
Diffstat (limited to 'src/core/ITensorPack.cpp')
| -rw-r--r-- | src/core/ITensorPack.cpp | 28 |
1 files changed, 23 insertions, 5 deletions
diff --git a/src/core/ITensorPack.cpp b/src/core/ITensorPack.cpp index 7a54a8bc6b..0f8b0824f8 100644 --- a/src/core/ITensorPack.cpp +++ b/src/core/ITensorPack.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2020 Arm Limited. + * Copyright (c) 2020-2021 Arm Limited. * * SPDX-License-Identifier: MIT * @@ -27,20 +27,33 @@ namespace arm_compute { +ITensorPack::ITensorPack(std::initializer_list<PackElement> l) : _pack() +{ + for (auto &e : l) + { + _pack[e.id] = e; + } +} + void ITensorPack::add_tensor(int id, ITensor *tensor) { - _pack[id] = PackElement(tensor); + _pack[id] = PackElement(id, tensor); } void ITensorPack::add_tensor(int id, const ITensor *tensor) { - _pack[id] = PackElement(tensor); + _pack[id] = PackElement(id, tensor); +} + +void ITensorPack::add_const_tensor(int id, const ITensor *tensor) +{ + add_tensor(id, tensor); } const ITensor *ITensorPack::get_const_tensor(int id) const { auto it = _pack.find(id); - if(it != _pack.end()) + if (it != _pack.end()) { return it->second.ctensor != nullptr ? it->second.ctensor : it->second.tensor; } @@ -53,6 +66,11 @@ ITensor *ITensorPack::get_tensor(int id) return it != _pack.end() ? it->second.tensor : nullptr; } +void ITensorPack::remove_tensor(int id) +{ + _pack.erase(id); +} + size_t ITensorPack::size() const { return _pack.size(); @@ -62,4 +80,4 @@ bool ITensorPack::empty() const { return _pack.empty(); } -} // namespace arm_compute
\ No newline at end of file +} // namespace arm_compute |