MLBEDSW-6941: Set correct OFM shape for fc op
If IFM operator shape is rewritten so that batching
is greater than one for fully connect, the OFM batch
must also be calculated. This change will fix output diffs
for networks that have fully connect OFM with rank greater
than 2.
Signed-off-by: Johan Alfven <johan.alfven@arm.com>
Change-Id: I5009edc647a1449a02c8116b45808c1c68beffe6
diff --git a/ethosu/vela/operation.py b/ethosu/vela/operation.py
index 8189793..4a56f1f 100644
--- a/ethosu/vela/operation.py
+++ b/ethosu/vela/operation.py
@@ -879,10 +879,8 @@
else:
# Special case, handled in graph optimization
self.ifm_shapes.append(Shape4D(ifm_tensor.get_full_shape()))
- if len(self.ofm.shape) == 2:
- self.ofm_shapes.append(Shape4D([self.ofm.shape[0], 1, 1, self.ofm.shape[1]]))
- else:
- self.ofm_shapes.append(Shape4D(ofm_tensor.get_full_shape()))
+ self.ofm_shapes.append(Shape4D(ofm_tensor.get_full_shape()))
+
elif self.type == Op.Softmax:
self.ifm_shapes.append(Shape4D(ifm_tensor.get_full_shape()))
self.ofm_shapes.append(Shape4D(ofm_tensor.get_full_shape()))