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(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( - 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);