IVGCVSW-2748 Fix bug causing std::bad_alloc in MakeTensor()

Change-Id: Iff6533396e594e2716e432223780d662d99d06ec
Signed-off-by: Aron Virginas-Tar <Aron.Virginas-Tar@arm.com>

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);