diff options
| -rw-r--r-- | src/core/NEON/kernels/winograd/winograd_gemm.cpp | 5 |
1 files changed, 2 insertions, 3 deletions
diff --git a/src/core/NEON/kernels/winograd/winograd_gemm.cpp b/src/core/NEON/kernels/winograd/winograd_gemm.cpp index fcfa635232..b45f6f55d9 100644 --- a/src/core/NEON/kernels/winograd/winograd_gemm.cpp +++ b/src/core/NEON/kernels/winograd/winograd_gemm.cpp @@ -34,11 +34,10 @@ Tensor4DShape WinogradGEMM<kr, kc, itr, itc>::Convolution<TOut, TIn>::get_output const PaddingType padding ) { - // TODO Accept different kernel sizes return Tensor4DShape { in_shape.n_batches, - (padding == PADDING_SAME) ? in_shape.n_rows : in_shape.n_rows - 2, - (padding == PADDING_SAME) ? in_shape.n_cols : in_shape.n_cols - 2, + (padding == PADDING_SAME) ? in_shape.n_rows : in_shape.n_rows - (kernel_rows - 2), + (padding == PADDING_SAME) ? in_shape.n_cols : in_shape.n_cols - (kernel_cols - 2), kernel_shape.n_output_channels, in_shape.ordering }; |