4 files changed, 12 insertions, 0 deletions
diff --git a/delegate/classic/src/Split.hpp b/delegate/classic/src/Split.hpp
index aaa610259f..57b7f8074e 100644
--- a/delegate/classic/src/Split.hpp
+++ b/delegate/classic/src/Split.hpp
@@ -107,6 +107,8 @@ TfLiteStatus VisitSplitOperator(DelegateData& delegateData,
     splitterDimSizes[splitDim] /= numSplits;
 
     armnn::SplitterDescriptor splitDescriptor(numSplits, inputDimSize);
+    splitDescriptor.SetAxis(axis);
+
     for (unsigned int j = 0; j < numSplits; ++j)
     {
         // Set the size of the views.
@@ -301,6 +303,8 @@ TfLiteStatus VisitSplitVOperator(DelegateData& delegateData,
     }
 
     armnn::SplitterDescriptor splitDescriptor(numSplits, inputDimSize);
+    splitDescriptor.SetAxis(axis);
+
     unsigned int accumSplit = 0;
     for (unsigned int j = 0; j < numSplits; ++j)
     {
diff --git a/delegate/classic/src/Unpack.hpp b/delegate/classic/src/Unpack.hpp
index 2cd32564c1..b3336ec990 100644
--- a/delegate/classic/src/Unpack.hpp
+++ b/delegate/classic/src/Unpack.hpp
@@ -96,6 +96,8 @@ TfLiteStatus VisitUnpackOperator(DelegateData& delegateData,
     unpackDimSizes[unpackAxis] /= unpackNum;
 
     armnn::SplitterDescriptor splitDesc(unpackNum, static_cast<unsigned int>(unpackDimSizes.size()));
+    splitDesc.SetAxis(unpackAxis);
+
     for (unsigned int j = 0; j < unpackNum; ++j)
     {
         // Set the size of the views.
diff --git a/delegate/opaque/src/Split.hpp b/delegate/opaque/src/Split.hpp
index 2dbfa602fb..199f46b126 100644
--- a/delegate/opaque/src/Split.hpp
+++ b/delegate/opaque/src/Split.hpp
@@ -134,6 +134,8 @@ TfLiteStatus VisitSplitOperator(DelegateData& delegateData,
     splitterDimSizes[splitDim] /= numSplits;
 
     armnn::SplitterDescriptor splitDescriptor(numSplits, inputDimSize);
+    splitDescriptor.SetAxis(axis);
+
     for (int j = 0; j < numSplits; ++j)
     {
         // Set the size of the views.
@@ -367,7 +369,9 @@ TfLiteStatus VisitSplitVOperator(DelegateData& delegateData,
     }
 
     armnn::SplitterDescriptor splitDescriptor(numSplits, inputDimSize);
+    splitDescriptor.SetAxis(axis);
     unsigned int accumSplit = 0;
+
     for (int j = 0; j < numSplits; ++j)
     {
         unsigned int splitSize = armnn::numeric_cast<unsigned int>(splitsTensorData[j]);
diff --git a/delegate/opaque/src/Unpack.hpp b/delegate/opaque/src/Unpack.hpp
index 0956d1688e..525529ff7b 100644
--- a/delegate/opaque/src/Unpack.hpp
+++ b/delegate/opaque/src/Unpack.hpp
@@ -96,6 +96,8 @@ TfLiteStatus VisitUnpackOperator(DelegateData& delegateData,
     unpackDimSizes[unpackAxis] /= unpackNum;
 
     armnn::SplitterDescriptor splitDesc(unpackNum, static_cast<unsigned int>(unpackDimSizes.size()));
+    splitDesc.SetAxis(unpackAxis);
+
     for (unsigned int j = 0; j < unpackNum; ++j)
     {
         // Set the size of the views.