diff options
Diffstat (limited to 'delegate/classic/src')
-rw-r--r-- | delegate/classic/src/Split.hpp | 4 | ||||
-rw-r--r-- | delegate/classic/src/Unpack.hpp | 2 |
2 files changed, 6 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. |