MLBEDSW-3018 Change to check for avoiding NHCWB16
Fixed issue with checking if axis corresponds to C-dim
Signed-off-by: Patrik Gustavsson <patrik.gustavsson@arm.com>
Change-Id: I72d9fd2c9fca642b5ab326324a63111b01c5de98
diff --git a/ethosu/vela/graph_optimiser.py b/ethosu/vela/graph_optimiser.py
index f6b03f6..e5fbc25 100644
--- a/ethosu/vela/graph_optimiser.py
+++ b/ethosu/vela/graph_optimiser.py
@@ -103,7 +103,7 @@
# multiple of 16. This as, it is only then the address offset for the ofm, for all operations, will be 16 byte
# aligned. For other values of axis the address offsets will be 16 byte aligned, as they are all based on c = 0
# and those addresses are always 16 byte aligned due to the NHCWB16 format.
- if axis == (len(tens.shape) - 1):
+ if axis == -1 or axis == (len(tens.shape) - 1):
for op in tens.ops:
if op.attrs["concat_start"] % 16 != 0:
tens.avoid_NHCWB16 = True