From 7f5e37aad7bf42ea3ef428d13a8c56efe98b3948 Mon Sep 17 00:00:00 2001 From: Haichen Shen Date: Mon, 11 May 2020 10:26:03 -0700 Subject: [PATCH] [Fix] Fix conv2d alter op for arm cpu (#5532) --- topi/python/topi/arm_cpu/conv2d_alter_op.py | 2 ++ 1 file changed, 2 insertions(+) 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 --- 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, -- 2.7.4