From e74e55418dbd40b91e8b32b74611870fac9f206e Mon Sep 17 00:00:00 2001
From: Aron Virginas-Tar <Aron.Virginas-Tar@arm.com>
Date: Thu, 21 Feb 2019 14:57:08 +0000
Subject: IVGCVSW-2748 Fix bug causing std::bad_alloc in MakeTensor()

Change-Id: Iff6533396e594e2716e432223780d662d99d06ec
Signed-off-by: Aron Virginas-Tar <Aron.Virginas-Tar@arm.com>
---
 src/armnn/test/TensorHelpers.hpp | 18 ++++++++++++++++--
 1 file 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);
-- 
cgit v1.2.1