diff options
Diffstat (limited to 'src/armnn/Layer.cpp')
-rw-r--r-- | src/armnn/Layer.cpp | 15 |
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()); } } |