diff --git a/onnxruntime/core/providers/cuda/math/topk_impl.cu b/onnxruntime/core/providers/cuda/math/topk_impl.cu
index 3dc4852695..466cdb0154 100644
--- a/onnxruntime/core/providers/cuda/math/topk_impl.cu
+++ b/onnxruntime/core/providers/cuda/math/topk_impl.cu
@@ -372,7 +372,7 @@ template <typename T>
 Status TopKImpl(const CudaKernel* kernel, const T* input_x, T* output_v, int64_t* output_i, const int64_t* elem_nums, size_t size, int64_t axis, int64_t K, int64_t largest, int64_t sorted, int64_t N, int64_t dimension) {
   auto aligned_K = ALIGN(K);
   auto aligned_dimension = ALIGN(dimension);
-  if (aligned_dimension <= GridDim::maxThreadsPerBlock << 1) {
+  if (aligned_dimension <= GridDim::maxThreadsPerBlock) {
     BitonicTopK<T><<<N, GridDim::maxThreadsPerBlock, aligned_dimension * sizeof(KV<T>)>>>(input_x, output_v, output_i, elem_nums, size, axis, K, aligned_K, largest, sorted, dimension, aligned_dimension, std::numeric_limits<T>::lowest(), std::numeric_limits<T>::max());
   } else if (K <= BT*16 || 0 == sorted) {
     auto XPT = static_cast<int64_t>(ceil(static_cast<double>(dimension) / GridDim::maxThreadsPerBlock));
diff --git a/onnxruntime/test/providers/cpu/math/topk_op_test.cc b/onnxruntime/test/providers/cpu/math/topk_op_test.cc
index 97b493ce8f..d3a01de6e3 100644
--- a/onnxruntime/test/providers/cpu/math/topk_op_test.cc
+++ b/onnxruntime/test/providers/cpu/math/topk_op_test.cc
@@ -475,8 +475,22 @@ TEST(TopKOperator, SortedSelection) {
   RunTest(11, 5, input_vals, input_dimensions, expected_vals, expected_indices, expected_dimensions, false, axis, 0);  // smallest values
 }
 
-TEST(TopKOperator, MediumArrayTopKSorted) 
-{
+// test dimension in range (GridDim::maxThreadsPerBlock, GridDim::maxThreadsPerBlock * 2], ie. [257, 512]
+TEST(TopKOperator, SmallArrayTopKSorted) {
+  std::vector<float> input_vals(400, 0.0f);
+  std::iota(input_vals.begin(), input_vals.end(), 0.0f);
+  std::vector<int64_t> input_dimensions = {400};
+  std::vector<float> expected_vals(400, 0.0f);
+  std::iota(expected_vals.begin(), expected_vals.end(), 0.0f);
+  std::reverse(expected_vals.begin(), expected_vals.end());
+  std::vector<int64_t> expected_indices(400, 0);
+  std::iota(expected_indices.begin(), expected_indices.end(), 0);
+  std::reverse(expected_indices.begin(), expected_indices.end());
+  std::vector<int64_t> expected_dimensions = {400};
+  RunTest(11, 400, input_vals, input_dimensions, expected_vals, expected_indices, expected_dimensions, false, -1, 1, 1);
+}
+
+TEST(TopKOperator, MediumArrayTopKSorted) {
   std::vector<float> input_vals(1000, 0.0f);
   std::iota(input_vals.begin(), input_vals.end(), 0.0f);
   std::vector<int64_t> input_dimensions = {1000};
@@ -490,8 +504,7 @@ TEST(TopKOperator, MediumArrayTopKSorted)
   RunTest(11, 100, input_vals, input_dimensions, expected_vals, expected_indices, expected_dimensions, false, 0, 1, 1);
 }
 
-TEST(TopKOperator, BigArrayTopKSorted) 
-{
+TEST(TopKOperator, BigArrayTopKSorted) {
   std::vector<float> input_vals(10000, 0.0f);
   std::iota(input_vals.begin(), input_vals.end(), 0.0f);
   std::vector<int64_t> input_dimensions = {10000};
@@ -505,8 +518,7 @@ TEST(TopKOperator, BigArrayTopKSorted)
   RunTest(11, 1000, input_vals, input_dimensions, expected_vals, expected_indices, expected_dimensions, false, 0, 1, 1);
 }
 
-TEST(TopKOperator, BigArrayBigTopKSorted) 
-{
+TEST(TopKOperator, BigArrayBigTopKSorted) {
   std::vector<float> input_vals(10000, 0.0f);
   std::iota(input_vals.begin(), input_vals.end(), 0.0f);
   std::vector<int64_t> input_dimensions = {10000};