commit da2b0030220e87788573a724979626aa92afd13e
author Patrik Gustavsson <patrik.gustavsson@arm.com> Thu Feb 04 16:28:29 2021 +0100
committer Patrik Gustavsson <patrik.gustavsson@arm.com> Mon Feb 08 14:24:41 2021 +0100
tree 66b0400f90e6854129005e08232f7db94c64fd14
parent 455e20e5ed0d5ce141a921e67f0219e55044e6e1

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

ethosu/vela/graph_optimiser.py

diff --git a/ethosu/vela/graph_optimiser.py b/ethosu/vela/graph_optimiser.py
index 1e3b131..5c1b90b 100644
--- a/ethosu/vela/graph_optimiser.py
+++ b/ethosu/vela/graph_optimiser.py
@@ -251,7 +251,6 @@
     if op.type == Op.Conv2DBackpropInput:
         # flip the inputs
         op.inputs[0], op.inputs[2] = op.inputs[2], op.inputs[0]
-        op.set_ifm_ofm_shapes()
         op.type = Op.Conv2DBackpropInputSwitchedBias
         op.ifm.resampling_mode = resampling_mode.TRANSPOSE
 
@@ -370,10 +369,9 @@
         batch_size = elms // n_in_elems
         assert batch_size * n_in_elems == elms
 
-        if op.ifm.shape != [batch_size, n_in_elems]:
-            op.ifm.avoid_NHCWB16 = True
-
         op.ifm_shapes[0] = Shape4D([batch_size, 1, 1, n_in_elems])
+        if Shape4D(op.ifm.shape) != op.ifm_shapes[0]:
+            op.ifm.avoid_NHCWB16 = True
     return op