commit 2e8a2dc31520213fa7e0cc096ad5d30ac010ba9a
author Georgios Pinitas <georgios.pinitas@arm.com> Fri Nov 09 17:59:00 2018 +0000
committer Georgios Pinitas <georgios.pinitas@arm.com> Mon Nov 12 12:22:14 2018 +0000
tree a5abe219d0a39ba16b8ceba6adb0a226c143b3db
parent e7f4a432775eb215c9da7ade1ee2710f67641068

COMPMID-1451: Set axis correctly in CLL2Normalize validate function

Change-Id: I93b14106cda8a1f640cf5acf120d31e2ebdaf495

src/runtime/CL/functions/CLL2NormalizeLayer.cpp

diff --git a/src/runtime/CL/functions/CLL2NormalizeLayer.cpp b/src/runtime/CL/functions/CLL2NormalizeLayer.cpp
index a3010a7..4f709d5 100644
--- a/src/runtime/CL/functions/CLL2NormalizeLayer.cpp
+++ b/src/runtime/CL/functions/CLL2NormalizeLayer.cpp
@@ -63,8 +63,8 @@
 
     ARM_COMPUTE_RETURN_ON_ERROR(CLReductionOperation::validate(input, &sum_sq, axis, ReductionOperation::SUM_SQUARE));
 
-    // Reduce shape on axis (supported axis is 0)
-    shape.set(0, 1);
+    // Reduce shape on axis
+    shape.set(axis, 1);
     sum_sq.set_tensor_shape(shape);
 
     ARM_COMPUTE_RETURN_ON_ERROR(CLL2NormalizeLayerKernel::validate(input, &sum_sq, output, axis, epsilon));