aboutsummaryrefslogtreecommitdiff
path: root/ethosu/vela/pass_packing.py
diff options
context:
space:
mode:
Diffstat (limited to 'ethosu/vela/pass_packing.py')
-rw-r--r--ethosu/vela/pass_packing.py4
1 files changed, 3 insertions, 1 deletions
diff --git a/ethosu/vela/pass_packing.py b/ethosu/vela/pass_packing.py
index 7abf3b24..ee0d7128 100644
--- a/ethosu/vela/pass_packing.py
+++ b/ethosu/vela/pass_packing.py
@@ -244,6 +244,7 @@ def pack_into_passes(nng, arch, verbose_packing=False):
input_set = set()
ifm_tensor = None
primary_op = None
+ ifm_shapes = None
to_process = collections.deque()
for start_op in start_ops_to_process:
@@ -279,6 +280,7 @@ def pack_into_passes(nng, arch, verbose_packing=False):
):
assert len(curr_op.inputs) >= 1
ifm_tensor = curr_op.ifm
+ ifm_shapes = curr_op.ifm_shapes.copy()
assert ifm_tensor is not None, "IFM missing in {}".format(curr_op)
assert ifm_tensor.purpose == TensorPurpose.FeatureMap
@@ -417,7 +419,7 @@ def pack_into_passes(nng, arch, verbose_packing=False):
ps.ifm_tensor = ifm_tensor
ps.ifm2_tensor = None
if ps.primary_op is not None and ps.primary_op.run_on_npu:
- ps.ifm_shapes.append(ps.primary_op.ifm_shapes[0])
+ ps.ifm_shapes.append(ifm_shapes[0])
ps.ofm_tensor = ofm_tensor
ps.ofm_shapes.append(ofm_shape)