diff options
| author | Dwight Lidman <dwight.lidman@arm.com> | 2021-05-03 13:12:57 +0200 |
|---|---|---|
| committer | Dwight Lidman <dwight.lidman@arm.com> | 2021-05-03 13:24:33 +0200 |
| commit | 8244ce0e5c2ca8e0bc788dea6fea524f7eae7b59 (patch) | |
| tree | d44da5ab97238952080f8d3d7bdcc58a04464ff8 | |
| parent | 5fabfcaa2b636b02899b4d6e0ccf95d853986475 (diff) | |
| download | ethos-u-vela-8244ce0e5c2ca8e0bc788dea6fea524f7eae7b59.tar.gz | |
MLBEDSW-4539: MEAN axis check exception fix
A recent fix to another MEAN bug introduced a new
bug. The bug was due to some incorrect logic for
checking the axis attribute.
Signed-off-by: Dwight Lidman <dwight.lidman@arm.com>
Change-Id: I65d3486a12e029f7c4450074f03fcd1974f65d8a
| -rw-r--r-- | ethosu/vela/graph_optimiser.py | 4 |
1 files changed, 2 insertions, 2 deletions
diff --git a/ethosu/vela/graph_optimiser.py b/ethosu/vela/graph_optimiser.py index 573b7ebb..99d4bf05 100644 --- a/ethosu/vela/graph_optimiser.py +++ b/ethosu/vela/graph_optimiser.py @@ -1476,8 +1476,8 @@ def convert_mean_to_depthwise_conv_or_avgpool(op, arch, nng): dims = len(shape) # Height and width axes have different index depending on dimensions - if len(axis.shape) <= 1: # single axis - axis = int(axis.values) if len(axis.shape) == 0 else axis.values[0] + if axis.shape == [] or axis.shape[0] == 1: # single axis + axis = int(axis.values) if len(axis.shape) == 0 else int(axis.values[0]) if dims in (2, 3): if axis == 0: h, w = shape[axis], 1 |