From c577f2c6a3b4ddb6ba87a882723c53a248afbeba Mon Sep 17 00:00:00 2001
From: telsoa01 <telmo.soares@arm.com>
Date: Fri, 31 Aug 2018 09:22:23 +0100
Subject: Release 18.08

---
 src/armnn/layers/ReshapeLayer.cpp | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

(limited to 'src/armnn/layers/ReshapeLayer.cpp')

diff --git a/src/armnn/layers/ReshapeLayer.cpp b/src/armnn/layers/ReshapeLayer.cpp
index df5d9d5bb0..248a45c491 100644
--- a/src/armnn/layers/ReshapeLayer.cpp
+++ b/src/armnn/layers/ReshapeLayer.cpp
@@ -30,17 +30,23 @@ ReshapeLayer* ReshapeLayer::Clone(Graph& graph) const
     return CloneBase<ReshapeLayer>(graph, m_Param, GetName());
 }
 
+std::vector<TensorShape> ReshapeLayer::InferOutputShapes(const std::vector<TensorShape>& inputShapes) const
+{
+    return std::vector<TensorShape>({ m_Param.m_TargetShape });
+}
+
 void ReshapeLayer::ValidateTensorShapesFromInputs()
 {
-    ConditionalThrow<LayerValidationException>(GetInputSlot(0).GetConnection() != nullptr,
-                     "ReshapeLayer: InputSlot must be connected to an OutputSlot");
-    ConditionalThrow<LayerValidationException>(GetInputSlot(0).GetConnection()->IsTensorInfoSet(),
-                     "ReshapeLayer: TensorInfo must be set on connected OutputSlot.");
+    VerifyLayerConnections(1, CHECK_LOCATION());
+
+    auto inferredShapes = InferOutputShapes({  });
+
+    BOOST_ASSERT(inferredShapes.size() == 1);
 
     ConditionalThrowIfNotEqual<LayerValidationException>(
         "ReshapeLayer: TensorShape set on OutputSlot[0] does not match the inferred shape.",
         GetOutputSlot(0).GetTensorInfo().GetShape(),
-        m_Param.m_TargetShape);
+        inferredShapes[0]);
 }
 
 } // namespace armnn
-- 
cgit v1.2.1