diff --git a/ethosu/vela/high_level_command_stream_generator.py b/ethosu/vela/high_level_command_stream_generator.py
index aad3783..c01790a 100644
--- a/ethosu/vela/high_level_command_stream_generator.py
+++ b/ethosu/vela/high_level_command_stream_generator.py
@@ -46,7 +46,12 @@
     npu_block_type = ps.npu_block_type
     split_offsets = list(ps.primary_op.read_offsets)  # offset for [ifm, ifm2]
 
-    if ps.ifm_tensor is not None and ps.ifm2_tensor is not None and npu_block_type == NpuBlockType.ElementWise:
+    if (
+        len(ps.inputs) == 2
+        and ps.ifm_tensor is not None
+        and ps.ifm2_tensor is not None
+        and npu_block_type == NpuBlockType.ElementWise
+    ):
         # Ensure correct ifm and ifm2 order
         if ps.inputs[0] == ps.primary_op.inputs[1] and ps.inputs[1] == ps.primary_op.inputs[0]:
             ps.ifm_tensor, ps.ifm2_tensor = ps.ifm2_tensor, ps.ifm_tensor