diff --git a/ethosu/regor/tflite/tflite_writer.cpp b/ethosu/regor/tflite/tflite_writer.cpp
index 79005e23113d977261bb129e18cdc65c4f19a86c..db70f242322aa2247647675c4069b6e6ac1d97b8 100644
--- a/ethosu/regor/tflite/tflite_writer.cpp
+++ b/ethosu/regor/tflite/tflite_writer.cpp
@@ -239,18 +239,30 @@ std::vector<const Tensor *> TfLiteWriter::SortedInputTensors(const Operation *op
 {
     std::vector<const Tensor *> tensors;
 
-    int ifm = 0;
-    for ( const auto &pair : TfLiteMapping::InputTensorIndices(type) )
+    const auto tensorIndices = TfLiteMapping::InputTensorIndices(type);
+    if ( tensorIndices.begin() != tensorIndices.end() )
     {
-        const TensorUsage usage = pair.second;
-        const auto conn = operation->Input(usage);
-        tensors.push_back(conn ? conn->tensor.get() : nullptr);
-        ifm += IsIFM(usage);
+        // If we have tensor indices for this op type, use that tensor order
+        int ifm = 0;
+        for ( const auto &[type_, usage] : tensorIndices )
+        {
+            const auto conn = operation->Input(usage);
+            tensors.push_back(conn ? conn->tensor.get() : nullptr);
+            ifm += IsIFM(usage);
+        }
+        while ( operation->Input(MakeTensorUsage(TensorUsage::IFM, ifm)) )
+        {
+            tensors.push_back(operation->IFM(ifm));
+            ifm++;
+        }
     }
-    while ( operation->Input(MakeTensorUsage(TensorUsage::IFM, ifm)) )
+    else
     {
-        tensors.push_back(operation->IFM(ifm));
-        ifm++;
+        // If we don't have tensor indices for this op type, use the tensor order we have
+        for ( const auto &[usage, conn] : operation->Inputs().pairs() )
+        {
+            tensors.push_back(conn.tensor.get());
+        }
     }
     return tensors;
 }