diff options
Diffstat (limited to 'ethosu/vela/pass_packing.py')
| -rw-r--r-- | ethosu/vela/pass_packing.py | 18 |
1 files changed, 10 insertions, 8 deletions
diff --git a/ethosu/vela/pass_packing.py b/ethosu/vela/pass_packing.py index 35f53300..7abf3b24 100644 --- a/ethosu/vela/pass_packing.py +++ b/ethosu/vela/pass_packing.py @@ -402,15 +402,17 @@ def pack_into_passes(nng, arch, verbose_packing=False): # 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 |