[Fix] Fix conv2d alter op for arm cpu (apache#5532)

trevor-m · May 22, 2020 · 17f1aad · 17f1aad
1 parent 6363edc
commit 17f1aad
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/topi/python/topi/arm_cpu/conv2d_alter_op.py b/topi/python/topi/arm_cpu/conv2d_alter_op.py
@@ -117,6 +117,7 @@ def _alter_conv2d_layout(attrs, inputs, tinfos, out_type):
         weight_expr = relay.transpose(weight_expr, axes=[0, 1, 2, 4, 3])
 
         new_attrs['tile_size'] = tile_size
+        new_attrs['channels'] = CO
 
         new_data = data
         new_kernel = te.placeholder((KH + tile_size - 1,
@@ -135,6 +136,7 @@ def _alter_conv2d_layout(attrs, inputs, tinfos, out_type):
         assert data_layout == "NCHW" and kernel_layout == "OIHW"
         N, CI, H, W = get_const_tuple(data.shape)
         CO, _, KH, KW = get_const_tuple(kernel.shape)
+        new_attrs['channels'] = CO
 
         # pre-compute winograd_nnpack transform
         # for winograd_nnpack_fp16, the the precompute prune pass must run on device,