From 02dcaec6cddbb05deba94ec06268a73784c36ac8 Mon Sep 17 00:00:00 2001
From: achirkin <achirkin@users.noreply.github.com>
Date: Thu, 10 Mar 2022 14:03:54 +0100
Subject: [PATCH] Use the pooled allocator helper

---
 cpp/bench/spatial/selection.cu | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/cpp/bench/spatial/selection.cu b/cpp/bench/spatial/selection.cu
index 4b15a00b7d..fbe65c2dbc 100644
--- a/cpp/bench/spatial/selection.cu
+++ b/cpp/bench/spatial/selection.cu
@@ -48,10 +48,7 @@ struct selection : public fixture {
 
   void run_benchmark(::benchmark::State& state) override
   {
-    rmm::mr::cuda_memory_resource cuda_mr;
-    rmm::mr::pool_memory_resource<rmm::mr::cuda_memory_resource> pool_mr{
-      &cuda_mr, size_t(1) << size_t(30), size_t(16) << size_t(30)};
-    rmm::mr::set_current_device_resource(&pool_mr);
+    using_pool_memory_res res;
     try {
       std::ostringstream label_stream;
       label_stream << params_.n_inputs << "#" << params_.input_len << "#" << params_.k;
@@ -71,7 +68,6 @@ struct selection : public fixture {
     } catch (raft::exception& e) {
       state.SkipWithError(e.what());
     }
-    rmm::mr::set_current_device_resource(nullptr);
   }
 
  private: