diff options
Diffstat (limited to 'src/armnn/layers/StridedSliceLayer.cpp')
-rw-r--r-- | src/armnn/layers/StridedSliceLayer.cpp | 14 |
1 files changed, 10 insertions, 4 deletions
diff --git a/src/armnn/layers/StridedSliceLayer.cpp b/src/armnn/layers/StridedSliceLayer.cpp index b9c337188c..dd4f942086 100644 --- a/src/armnn/layers/StridedSliceLayer.cpp +++ b/src/armnn/layers/StridedSliceLayer.cpp @@ -52,15 +52,21 @@ std::vector<TensorShape> StridedSliceLayer::InferOutputShapes( for (unsigned int i = 0; i < inputShape.GetNumDimensions(); i++) { + int stride = m_Param.m_Stride[i]; + int start = m_Param.GetStartForAxis(inputShape, i); + int stop = m_Param.GetStopForAxis(inputShape, i, start); + if (m_Param.m_ShrinkAxisMask & (1 << i)) { + // Don't take a slice from an axis being shrunk + if (m_Param.m_End[i] >= 2) + { + throw LayerValidationException( + "StridedSlice: Attempting to take slice from an axis being shrunk"); + } continue; } - int stride = m_Param.m_Stride[i]; - int start = m_Param.GetStartForAxis(inputShape, i); - int stop = m_Param.GetStopForAxis(inputShape, i, start); - int newSize = stride > 0 ? ((stop - start) + stride - 1) / stride : ((start - stop) - stride - 1) / -stride; |