MLBEDSW-3804 Add missing check for npu_op
Add missing check for npu_op
Op ifm ofm shapes only valid for npu_ops.
Signed-off-by: Patrik Gustavsson <patrik.gustavsson@arm.com>
Change-Id: I73624c8e122fee510ab8320172b8b3a648a6f070
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