diff options
Diffstat (limited to 'src/armnn/Network.cpp')
-rw-r--r-- | src/armnn/Network.cpp | 16 |
1 files changed, 16 insertions, 0 deletions
diff --git a/src/armnn/Network.cpp b/src/armnn/Network.cpp index e00dbfc0fc..17a1da1f6c 100644 --- a/src/armnn/Network.cpp +++ b/src/armnn/Network.cpp @@ -934,6 +934,11 @@ OptimizationResult AssignBackends(OptimizedNetworkImpl* optNetObjPtr, { auto layer = *it; + if (layer->GetType() == LayerType::Input) + { + continue; + } + DataType dataTypeIn = layer->GetNumInputSlots() == 0 ? DataType::Float32 : layer->GetInputSlot(0).GetConnectedOutputSlot()->GetTensorInfo().GetDataType(); DataType dataTypeOut = layer->GetNumOutputSlots() == 0 ? DataType::Float32 : @@ -1027,6 +1032,17 @@ OptimizationResult AssignBackends(OptimizedNetworkImpl* optNetObjPtr, } } + for (auto it = firstLayer; it != lastLayer; ++it) + { + auto layer = *it; + + if(layer->GetType() == LayerType::Input) + { + BackendId connectedBackendId = layer->GetOutputSlot(0).GetConnection(0)->GetOwningLayer().GetBackendId(); + layer->SetBackendId(connectedBackendId); + } + } + return result; } |