diff options
Diffstat (limited to 'src/armnn/test')
-rw-r--r-- | src/armnn/test/CreateWorkload.hpp | 9 |
1 files changed, 7 insertions, 2 deletions
diff --git a/src/armnn/test/CreateWorkload.hpp b/src/armnn/test/CreateWorkload.hpp index fed9dd6b33..b9735f4f2f 100644 --- a/src/armnn/test/CreateWorkload.hpp +++ b/src/armnn/test/CreateWorkload.hpp @@ -885,9 +885,14 @@ std::unique_ptr<L2NormalizationWorkload> CreateL2NormalizationWorkloadTest(armnn Layer* const input = graph.AddLayer<InputLayer>(0, "input"); Layer* const output = graph.AddLayer<OutputLayer>(0, "output"); + TensorShape inputShape = (dataLayout == DataLayout::NCHW) ? + TensorShape{ 5, 20, 50, 67 } : TensorShape{ 5, 50, 67, 20 }; + TensorShape outputShape = (dataLayout == DataLayout::NCHW) ? + TensorShape{ 5, 20, 50, 67 } : TensorShape{ 5, 50, 67, 20 }; + // Connects up. - armnn::TensorInfo inputTensorInfo({ 5, 20, 50, 67 }, DataType); - armnn::TensorInfo outputTensorInfo({ 5, 20, 50, 67 }, DataType); + armnn::TensorInfo inputTensorInfo(inputShape, DataType); + armnn::TensorInfo outputTensorInfo(outputShape, DataType); Connect(input, layer, inputTensorInfo); Connect(layer, output, outputTensorInfo); CreateTensorHandles(graph, factory); |