From 8244ce0e5c2ca8e0bc788dea6fea524f7eae7b59 Mon Sep 17 00:00:00 2001 From: Dwight Lidman Date: Mon, 3 May 2021 13:12:57 +0200 Subject: 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 Change-Id: I65d3486a12e029f7c4450074f03fcd1974f65d8a --- ethosu/vela/graph_optimiser.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) (limited to 'ethosu/vela/graph_optimiser.py') 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 -- cgit v1.2.1