From 835d8e10f33f411664cebe65d3f6a872f6cc849a Mon Sep 17 00:00:00 2001 From: Fredrik Svedberg Date: Fri, 4 Sep 2020 09:46:17 +0200 Subject: [MLBEDSW-2928] Add batching to softmax Added batching to softmax by reshaping the input. Signed-off-by: Fredrik Svedberg Change-Id: I0b516f9bf2410fb86372b229beba4a7280c498cc --- ethosu/vela/supported_operators.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) (limited to 'ethosu/vela/supported_operators.py') diff --git a/ethosu/vela/supported_operators.py b/ethosu/vela/supported_operators.py index e0ee6163..86cc3c07 100644 --- a/ethosu/vela/supported_operators.py +++ b/ethosu/vela/supported_operators.py @@ -420,8 +420,8 @@ class SupportedOperators: if ifm_tensor.dtype not in (DataType.uint8, DataType.int8, DataType.int16): return False - # check batch size - if len(ifm_tensor.shape) in (2, 4) and ifm_tensor.shape[0] != 1: + # check shape + if len(ifm_tensor.shape) > 4 or ifm_tensor.shape != ofm_tensor.shape: return False return True -- cgit v1.2.1