Release 18.08

[platform/upstream/armnn.git] / src / armnn / layers / ActivationLayer.cpp

diff --git a/src/armnn/layers/ActivationLayer.cpp b/src/armnn/layers/ActivationLayer.cpp
index 2371eaa..ad1e4a9 100644 (file)
--- a/src/armnn/layers/ActivationLayer.cpp
+++ b/src/armnn/layers/ActivationLayer.cpp
@@ -30,12 +30,16 @@ ActivationLayer* ActivationLayer::Clone(Graph& graph) const
 
 void ActivationLayer::ValidateTensorShapesFromInputs()
 {
-    auto& info = GetInputSlot(0).GetConnection()->GetTensorInfo();
+    VerifyLayerConnections(1, CHECK_LOCATION());
+
+    auto inferredShapes = InferOutputShapes({ GetInputSlot(0).GetConnection()->GetTensorInfo().GetShape() });
+
+    BOOST_ASSERT(inferredShapes.size() == 1);
 
     ConditionalThrowIfNotEqual<LayerValidationException>(
         "ActivationLayer: TensorShape set on OutputSlot[0] does not match the inferred shape.",
         GetOutputSlot(0).GetTensorInfo().GetShape(),
-        info.GetShape());
+        inferredShapes[0]);
 }
 
 } // namespace armnn