diff options
Diffstat (limited to 'src/armnn/layers/Pooling2dLayer.cpp')
| -rw-r--r-- | src/armnn/layers/Pooling2dLayer.cpp | 8 |
1 files changed, 5 insertions, 3 deletions
diff --git a/src/armnn/layers/Pooling2dLayer.cpp b/src/armnn/layers/Pooling2dLayer.cpp index 821c011cd0..24b7a69c49 100644 --- a/src/armnn/layers/Pooling2dLayer.cpp +++ b/src/armnn/layers/Pooling2dLayer.cpp @@ -7,6 +7,7 @@ #include "LayerCloneBase.hpp" #include <armnn/TypesUtils.hpp> +#include <backendsCommon/DataLayoutIndexed.hpp> #include <backendsCommon/WorkloadData.hpp> #include <backendsCommon/WorkloadFactory.hpp> @@ -33,13 +34,14 @@ std::vector<TensorShape> Pooling2dLayer::InferOutputShapes(const std::vector<Ten { BOOST_ASSERT(inputShapes.size() == 1); const TensorShape& inputShape = inputShapes[0]; + const DataLayoutIndexed dimensionIndices = m_Param.m_DataLayout; // If we support multiple batch dimensions in the future, then this assert will need to change. BOOST_ASSERT_MSG(inputShape.GetNumDimensions() == 4, "Pooling2dLayer will always have 4D input."); - unsigned int inWidth = inputShape[m_Param.m_DataLayout.GetWidthIndex()]; - unsigned int inHeight = inputShape[m_Param.m_DataLayout.GetHeightIndex()]; - unsigned int inChannels = inputShape[m_Param.m_DataLayout.GetChannelsIndex()]; + unsigned int inWidth = inputShape[dimensionIndices.GetWidthIndex()]; + unsigned int inHeight = inputShape[dimensionIndices.GetHeightIndex()]; + unsigned int inChannels = inputShape[dimensionIndices.GetChannelsIndex()]; unsigned int inBatchSize = inputShape[0]; bool isGlobalPooling = (m_Param.m_StrideX==0 && m_Param.m_StrideY==0); |