diff --git a/ethosu/vela/operation.py b/ethosu/vela/operation.py
index 3b34fe8..c1ca3f8 100644
--- a/ethosu/vela/operation.py
+++ b/ethosu/vela/operation.py
@@ -312,6 +312,15 @@
         if self not in tens.consumer_list:
             tens.consumer_list.append(self)
 
+    def set_input_tensor(self, tens, idx):
+        tens_to_remove = self.inputs[idx]
+        if tens_to_remove in tens.consumer_list:
+            tens.consumer_list.remove(tens_to_remove)
+
+        self.inputs[idx] = tens
+        if self not in tens.consumer_list:
+            tens.consumer_list.append(self)
+
     def set_output_tensor(self, tens):
         tens.ops = [self]
         self.outputs = [tens]