diff --git a/topi/python/topi/arm_cpu/conv2d_alter_op.py b/topi/python/topi/arm_cpu/conv2d_alter_op.py index aa4c878b437e..221ccce1a2a0 100644 --- 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,