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

diff --git a/topi/python/topi/arm_cpu/conv2d_alter_op.py b/topi/python/topi/arm_cpu/conv2d_alter_op.py
index aa4c878..221ccce 100644 (file)
--- a/topi/python/topi/arm_cpu/conv2d_alter_op.py
+++ b/topi/python/topi/arm_cpu/conv2d_alter_op.py
@@ -121,6 +121,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,
@@ -139,6 +140,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,