 
 template <typename T>
 bool SortTopk(const platform::CUDADeviceContext& ctx,
-              const framework::Tensor* input_tensor, const size_t num_cols,
-              const size_t num_rows, size_t k, framework::Tensor* out_tensor,
+              const framework::Tensor* input_tensor, const int64_t num_cols,
+              const int64_t num_rows, const int k,
+              framework::Tensor* out_tensor,
               framework::Tensor* indices_tensor) {
   auto cu_stream = ctx.stream();
 
   Tensor input_indices;
-  const std::vector<int64_t> dims = {static_cast<int64_t>(num_rows),
-                                     static_cast<int64_t>(num_cols)};
+  const std::vector<int64_t> dims = {num_rows, num_cols};
   auto dim = framework::make_ddim(dims);
   input_indices.Resize(dim);
   // input_indices.Resize(num_rows*num_cols);

  int num_cols_;
};

__global__ void InitIndex(int64_t* indices, int num_rows, int num_cols) {
  int col_id = threadIdx.x;
  int row_id = blockIdx.x;

  for (int j = row_id; j < num_rows; j += gridDim.x) {
    for (int i = col_id; i < num_cols; i += blockDim.x) {
      indices[j * num_cols + i] = i;
    }
  }

    // FIXME(typhoonzero): data is always converted to type T?

    framework::DDim inputdims = input->dims();
    const int64_t input_height = framework::product(
        framework::slice_ddim(inputdims, 0, inputdims.size() - 1));
    const int64_t input_width = inputdims[inputdims.size() - 1];
    const auto& dev_ctx = ctx.cuda_device_context();

    if ((input_width <= 1024 || k >= 128 || k == input_width)) {
