fix mini topk bug

jack603047588 · May 26, 2022 · 4f4d42d · 4f4d42d
1 parent 51acec7
commit 4f4d42d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/paddle/fluid/operators/top_k_op.cu b/paddle/fluid/operators/top_k_op.cu
@@ -66,7 +66,7 @@ __global__ void KernelSortTopK(const size_t num_rows, const T* input_val,
       }
     } else {
       for (int i = 0; i < K; ++i) {
-        int pos = (i + 1) % K;
+        int pos = (i + 1) % num_cols;
         val[i] = in[pos];
         ind[i] = pos;
       }