MLBEDSW-3937 Fix check for NHCWB16 for FC
Fix check for NHCWB16 for modifying FC input.
Signed-off-by: Patrik Gustavsson <patrik.gustavsson@arm.com>
Change-Id: Ie50c32ca079afadd0af9b7b909820794ceee373c
diff --git a/ethosu/vela/operation.py b/ethosu/vela/operation.py
index 09371b7..b297bed 100644
--- a/ethosu/vela/operation.py
+++ b/ethosu/vela/operation.py
@@ -739,7 +739,7 @@
if self.type == Op.Softmax:
self.ifm_shapes.append(Shape4D(ifm_tensor.get_full_shape()))
self.ofm_shapes.append(Shape4D(ofm_tensor.get_full_shape()))
- elif self.type.is_split_op or self.type.is_concat_op():
+ elif self.type.is_split_op() or self.type.is_concat_op():
for inp in self.inputs:
if inp is not None:
self.ifm_shapes.append(Shape4D(full_shape(4, inp.shape, 1)))
@@ -751,7 +751,9 @@
else:
self.ofm_shapes.append(None)
else:
- self.ifm_shapes.append(Shape4D(full_shape(4, ifm_tensor.shape, 1)))
+ if ifm_tensor is not None:
+ self.ifm_shapes.append(Shape4D(full_shape(4, ifm_tensor.shape, 1)))
if ifm2_tensor is not None:
self.ifm_shapes.append(Shape4D(full_shape(4, ifm2_tensor.shape, 1)))
- self.ofm_shapes.append(Shape4D(full_shape(4, ofm_tensor.shape, 1)))
+ if ofm_tensor is not None:
+ self.ofm_shapes.append(Shape4D(full_shape(4, ofm_tensor.shape, 1)))