From 3b686b189e81f57455abb6737b581d306987bbae Mon Sep 17 00:00:00 2001 From: Zhang Zheng <32410583+ZzSean@users.noreply.github.com> Date: Sat, 2 Apr 2022 10:30:54 +0800 Subject: [PATCH] Limit the condition of entering optimized kernel (#41296) Co-authored-by: root --- paddle/phi/kernels/gpu/top_k_kernel.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/paddle/phi/kernels/gpu/top_k_kernel.cu b/paddle/phi/kernels/gpu/top_k_kernel.cu index adaf5cc092b..8262023826b 100644 --- a/paddle/phi/kernels/gpu/top_k_kernel.cu +++ b/paddle/phi/kernels/gpu/top_k_kernel.cu @@ -98,7 +98,7 @@ void TopkKernel(const Context& dev_ctx, } #if defined(PADDLE_WITH_CUDA) && CUDA_VERSION >= 9000 - if (input_width >= 1024 && input_height == 1) { + if (input_width >= 1024 && in_dims.size() == 1) { // 1. Gather TopK, but without sorting constexpr int max_num_threads = 1024; if (largest) { -- GitLab