From 47a4385eb84dbc27639e8d65f6370f604bf642ba Mon Sep 17 00:00:00 2001
From: JiayuSun <470323747@qq.com>
Date: Fri, 28 Jul 2023 11:37:06 +0800
Subject: [PATCH] Fix missing parameter for select_k (#1682)

The parameter `rmm::mr::device_memory_resource* mr` is missing while passing down.

Authors:
  - JiayuSun (https://github.com/ucassjy)

Approvers:
  - Ben Frederickson (https://github.com/benfred)

URL: https://github.com/rapidsai/raft/pull/1682
---
 cpp/include/raft/matrix/detail/select_k-inl.cuh | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/cpp/include/raft/matrix/detail/select_k-inl.cuh b/cpp/include/raft/matrix/detail/select_k-inl.cuh
index b852f26e2e..af5a5770fb 100644
--- a/cpp/include/raft/matrix/detail/select_k-inl.cuh
+++ b/cpp/include/raft/matrix/detail/select_k-inl.cuh
@@ -303,7 +303,8 @@ void select_k(raft::resources const& handle,
                                                         out_idx,
                                                         select_min,
                                                         true,  // fused_last_filter
-                                                        stream);
+                                                        stream,
+                                                        mr);
 
       if (sorted) {
         auto offsets = raft::make_device_vector<IdxT, IdxT>(handle, (IdxT)(batch_size + 1));
@@ -326,7 +327,7 @@ void select_k(raft::resources const& handle,
     case Algo::kWarpDistributedShm:
       return detail::select::warpsort::
         select_k_impl<T, IdxT, detail::select::warpsort::warp_sort_distributed_ext>(
-          in_val, in_idx, batch_size, len, k, out_val, out_idx, select_min, stream);
+          in_val, in_idx, batch_size, len, k, out_val, out_idx, select_min, stream, mr);
     case Algo::kFaissBlockSelect:
       return neighbors::detail::select_k(
         in_val, in_idx, batch_size, len, out_val, out_idx, select_min, k, stream);