diff --git a/ethosu/vela/pass_packing.py b/ethosu/vela/pass_packing.py
index 35f5330..7abf3b2 100644
--- a/ethosu/vela/pass_packing.py
+++ b/ethosu/vela/pass_packing.py
@@ -402,15 +402,17 @@
 
             # Get the corresponding ifm_shapes
             for op in input_ops_list + [primary_op]:
-                if ps.ifm_tensor == op.ifm:
-                    ps.ifm_shapes.append(op.ifm_shapes[0])
-                elif ps.ifm_tensor == op.ifm2:
-                    ps.ifm_shapes.append(op.ifm_shapes[1])
+                if op.run_on_npu:
+                    if ps.ifm_tensor == op.ifm:
+                        ps.ifm_shapes.append(op.ifm_shapes[0])
+                    elif ps.ifm_tensor == op.ifm2:
+                        ps.ifm_shapes.append(op.ifm_shapes[1])
             for op in input_ops_list + [primary_op]:
-                if ps.ifm2_tensor == op.ifm:
-                    ps.ifm_shapes.append(op.ifm_shapes[0])
-                elif ps.ifm2_tensor == op.ifm2:
-                    ps.ifm_shapes.append(op.ifm_shapes[1])
+                if op.run_on_npu:
+                    if ps.ifm2_tensor == op.ifm:
+                        ps.ifm_shapes.append(op.ifm_shapes[0])
+                    elif ps.ifm2_tensor == op.ifm2:
+                        ps.ifm_shapes.append(op.ifm_shapes[1])
         else:
             ps.ifm_tensor = ifm_tensor
             ps.ifm2_tensor = None