From e0ee7403e77c5ee0a50359b12ef989a1ed8a2e7d Mon Sep 17 00:00:00 2001
From: Thomas Young <35565423+HexToString@users.noreply.github.com>
Date: Fri, 6 Jan 2023 11:02:24 +0800
Subject: [PATCH] fix bug (#49546)

---
 paddle/phi/kernels/cpu/top_k_kernel.cc | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/paddle/phi/kernels/cpu/top_k_kernel.cc b/paddle/phi/kernels/cpu/top_k_kernel.cc
index 3e946803660..cebe292d339 100644
--- a/paddle/phi/kernels/cpu/top_k_kernel.cc
+++ b/paddle/phi/kernels/cpu/top_k_kernel.cc
@@ -31,6 +31,14 @@ static void FullTopK(Type input_height,
                      const int& k,
                      const bool& largest,
                      const bool& sorted) {
+  PADDLE_ENFORCE_LE(
+      k,
+      input_width,
+      errors::InvalidArgument("The rank (%d) of the input 'k' for "
+                              "topk op must be less than or equal to %d.",
+                              k,
+                              input_width));
+
   // when the k is small, will the partial sort
   bool partial_sort_flag = (k * 64) < input_width;
 
-- 
GitLab