1 files changed, 14 insertions, 1 deletions
diff --git a/src/armnn/Layer.cpp b/src/armnn/Layer.cpp
index 8d4811ae89..d2f8f2c982 100644
--- a/src/armnn/Layer.cpp
+++ b/src/armnn/Layer.cpp
@@ -259,7 +259,20 @@ void Layer::CollectWorkloadInputs(WorkloadDataCollector& dataCollector) const
         // The graph must be well-formed at this point.
         ARMNN_ASSERT(inputSlot.GetConnection());
         const OutputHandler& outputHandler = inputSlot.GetConnectedOutputSlot()->GetOutputHandler();
-        dataCollector.Push(outputHandler.GetData(), outputHandler.GetTensorInfo());
+
+        if (inputSlot.IsTensorInfoOverridden() && outputHandler.GetData())
+        {
+            auto handler = outputHandler.GetData()->DecorateTensorHandle(inputSlot.GetTensorInfo());
+
+            if (handler)
+            {
+                // Add overridden TensorHandle
+                dataCollector.Push(handler.get(), inputSlot.GetTensorInfo());
+                continue;
+            }
+        }
+        // Add default TensorHandle
+        dataCollector.Push(outputHandler.GetData(), inputSlot.GetTensorInfo());
     }
 }