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;
 }