1 files changed, 30 insertions, 1 deletions
diff --git a/src/armnn/layers/PadLayer.cpp b/src/armnn/layers/PadLayer.cpp
index 4fcbc77c7b..a8c749c570 100644
--- a/src/armnn/layers/PadLayer.cpp
+++ b/src/armnn/layers/PadLayer.cpp
@@ -36,10 +36,39 @@ PadLayer* PadLayer::Clone(Graph& graph) const
     return std::move(layer);
 }
 
+std::vector<TensorShape> PadLayer::InferOutputShapes(const std::vector<TensorShape>& inputShapes) const
+{
+    ARMNN_ASSERT(inputShapes.size() == 1);
+    const TensorShape& inputShape = inputShapes[0];
+
+    unsigned int rank = inputShape.GetNumDimensions();
+    ARMNN_ASSERT(m_Param.m_PadList.size() == rank);
+    ARMNN_ASSERT(rank != 0);
+
+    std::vector<unsigned int> outputDimensionSizes;
+    outputDimensionSizes.reserve(rank);
+    for (unsigned int i = 0; i < rank; ++i)
+    {
+        outputDimensionSizes[i] = inputShape[i] + m_Param.m_PadList[i].first + m_Param.m_PadList[i].second;
+    }
+
+    TensorShape tensorShape = TensorShape( rank, outputDimensionSizes.data());
+    return std::vector<TensorShape>({ tensorShape });
+}
+
 void PadLayer::ValidateTensorShapesFromInputs()
 {
+    VerifyLayerConnections(1, CHECK_LOCATION());
+
+    const TensorShape& outputShape = GetOutputSlot(0).GetTensorInfo().GetShape();
+
+    VerifyShapeInferenceType(outputShape, m_ShapeInferenceMethod);
+
+    auto inferredShapes = InferOutputShapes({ GetInputSlot(0).GetConnection()->GetTensorInfo().GetShape() });
+
+    ARMNN_ASSERT(inferredShapes.size() == 1);
 
-    return;
+    ValidateAndCopyShape(outputShape, inferredShapes[0], m_ShapeInferenceMethod, "PadLayer");
 }
 
 void PadLayer::Accept(ILayerVisitor& visitor) const