1 files changed, 12 insertions, 4 deletions
diff --git a/src/backends/aclCommon/ArmComputeUtils.hpp b/src/backends/aclCommon/ArmComputeUtils.hpp
index f466ab1777..5d424afa9a 100644
--- a/src/backends/aclCommon/ArmComputeUtils.hpp
+++ b/src/backends/aclCommon/ArmComputeUtils.hpp
@@ -9,6 +9,7 @@
 #include <armnn/utility/Assert.hpp>
 #include <armnn/utility/NumericCast.hpp>
 #include <armnn/backends/WorkloadData.hpp>
+#include <armnnUtils/TensorUtils.hpp>
 
 #include <arm_compute/runtime/FunctionDescriptors.h>
 #include <arm_compute/function_info/FullyConnectedLayerInfo.h>
@@ -248,13 +249,20 @@ inline std::set<unsigned int> ComputeSplitAxis(const armnn::SplitterDescriptor&
     unsigned int numDimensions = desc.GetNumDimensions();
     std::set<unsigned int> splitAxis;
 
-    for (unsigned int i = 0; i < numSplit; ++i)
+    if (desc.HasAxis())
     {
-        for (unsigned int dimIdx = 0; dimIdx < numDimensions; ++dimIdx)
+        splitAxis.insert(armnnUtils::GetUnsignedAxis(desc.GetNumDimensions(), desc.GetAxis()));
+    }
+    else
+    {
+        for (unsigned int i = 0; i < numSplit; ++i)
         {
-            if (desc.GetViewSizes(i)[dimIdx] != input[dimIdx])
+            for (unsigned int dimIdx = 0; dimIdx < numDimensions; ++dimIdx)
             {
-                splitAxis.insert(dimIdx);
+                if (desc.GetViewSizes(i)[dimIdx] != input[dimIdx])
+                {
+                    splitAxis.insert(dimIdx);
+                }
             }
         }
     }