diff options
author | Aron Virginas-Tar <Aron.Virginas-Tar@arm.com> | 2019-02-21 14:57:08 +0000 |
---|---|---|
committer | Aron Virginas-Tar <aron.virginas-tar@arm.com> | 2019-02-21 16:38:24 +0000 |
commit | e74e55418dbd40b91e8b32b74611870fac9f206e (patch) | |
tree | b342c42e38601b3ac43d73116ca168d0012ea112 | |
parent | dbe94290bdd53b830604ae58f1f1d7253d6069f8 (diff) | |
download | armnn-e74e55418dbd40b91e8b32b74611870fac9f206e.tar.gz |
IVGCVSW-2748 Fix bug causing std::bad_alloc in MakeTensor()
Change-Id: Iff6533396e594e2716e432223780d662d99d06ec
Signed-off-by: Aron Virginas-Tar <Aron.Virginas-Tar@arm.com>
-rw-r--r-- | src/armnn/test/TensorHelpers.hpp | 18 |
1 files changed, 16 insertions, 2 deletions
diff --git a/src/armnn/test/TensorHelpers.hpp b/src/armnn/test/TensorHelpers.hpp index fcaa0772a0..35e471e0f2 100644 --- a/src/armnn/test/TensorHelpers.hpp +++ b/src/armnn/test/TensorHelpers.hpp @@ -198,9 +198,23 @@ boost::multi_array<T, n> MakeTensor(const armnn::TensorInfo& tensorInfo, const s std::array<unsigned int, n> shape; - for (unsigned int i = 0; i < n; i++) + // NOTE: tensorInfo.GetNumDimensions() might be different from n + const unsigned int returnDimensions = static_cast<unsigned int>(n); + const unsigned int actualDimensions = tensorInfo.GetNumDimensions(); + + const unsigned int paddedDimensions = + returnDimensions > actualDimensions ? returnDimensions - actualDimensions : 0u; + + for (unsigned int i = 0u; i < returnDimensions; i++) { - shape[i] = tensorInfo.GetShape()[i]; + if (i < paddedDimensions) + { + shape[i] = 1u; + } + else + { + shape[i] = tensorInfo.GetShape()[i - paddedDimensions]; + } } boost::const_multi_array_ref<T, n> arrayRef(&flat[0], shape); |