IVGCVSW-2051 Fixed the dynamic allocation of the dimSizes array in MeanLayer

Change-Id: I6e39f1e1bc32c0c2c1953da00016edd68c48d4e0

1 files changed, 3 insertions, 5 deletions

diff --git a/src/armnn/layers/MeanLayer.cpp b/src/armnn/layers/MeanLayer.cpp
index 3259d95afd..d1c4591c92 100644
--- a/src/armnn/layers/MeanLayer.cpp
+++ b/src/armnn/layers/MeanLayer.cpp
@@ -73,9 +73,7 @@ void MeanLayer::ValidateTensorShapesFromInputs()
         }
     }
 
-    unsigned int dimSizes[outputRank];
-    memset(dimSizes, 1, outputRank * sizeof(unsigned int));
-
+    std::vector<unsigned int> dimSizes(outputRank, 1);
     if (!m_Param.m_Axis.empty())
     {
         // Skip the dimension that has been reduced unless keepDims is true.
@@ -94,7 +92,7 @@ void MeanLayer::ValidateTensorShapesFromInputs()
             }
         }
     }
-    const TensorShape& inferredShape = TensorShape(outputRank, dimSizes);
+    const TensorShape& inferredShape = TensorShape(outputRank, dimSizes.data());
 
     ConditionalThrowIfNotEqual<LayerValidationException>(
         "MeanLayer: TensorShape set on OutputSlot[0] does not match the inferred shape.",
@@ -102,4 +100,4 @@ void MeanLayer::ValidateTensorShapesFromInputs()
         inferredShape);
 }
 
-} // namespace armnn
\ No newline at end of file
+} // namespace armnn