Fix reshape usage in ARM Winograd (apache#5732)

trevor-m · Jun 9, 2020 · 6996fb6 · 6996fb6
1 parent 50d79ff
commit 6996fb6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/topi/python/topi/arm_cpu/conv2d_alter_op.py b/topi/python/topi/arm_cpu/conv2d_alter_op.py
@@ -117,7 +117,7 @@ def _alter_conv2d_layout(attrs, inputs, tinfos, out_type):
         weight_expr = relay.reshape(weight_expr,
                                     newshape=(KH + tile_size - 1,
                                               KW + tile_size - 1,
-                                              idxd(CO, VC), VC, CI))
+                                              CO // VC, VC, CI))
         weight_expr = relay.transpose(weight_expr, axes=[0, 1, 2, 4, 3])
 
         new_attrs['tile_size'] = tile_size