1 files changed, 9 insertions, 0 deletions
diff --git a/ethosu/vela/operation.py b/ethosu/vela/operation.py
index 3b34fe81..c1ca3f81 100644
--- a/ethosu/vela/operation.py
+++ b/ethosu/vela/operation.py
@@ -312,6 +312,15 @@ input and output tensors, as well as an attribute dictionary."""
         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]