IVGCVSW-2748 Fix bug causing std::bad_alloc in MakeTensor()
Change-Id: Iff6533396e594e2716e432223780d662d99d06ec
Signed-off-by: Aron Virginas-Tar <Aron.Virginas-Tar@arm.com>
diff --git a/src/armnn/test/TensorHelpers.hpp b/src/armnn/test/TensorHelpers.hpp
index fcaa077..35e471e 100644
--- a/src/armnn/test/TensorHelpers.hpp
+++ b/src/armnn/test/TensorHelpers.hpp
@@ -198,9 +198,23 @@
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);