From 629e6fa3c50ea149701ebc22e17c820402c02831 Mon Sep 17 00:00:00 2001
From: Tamas Bela Feher <tfeher@nvidia.com>
Date: Thu, 18 Jan 2024 13:35:16 +0100
Subject: [PATCH] Use managed memory for the new temprorary buffer

---
 cpp/include/raft/neighbors/detail/ivf_pq_build.cuh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cpp/include/raft/neighbors/detail/ivf_pq_build.cuh b/cpp/include/raft/neighbors/detail/ivf_pq_build.cuh
index 65278d7944..d2105cd354 100644
--- a/cpp/include/raft/neighbors/detail/ivf_pq_build.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_pq_build.cuh
@@ -1763,7 +1763,7 @@ auto build(raft::resources const& handle,
     // Besides just sampling, we transform the input dataset into floats to make it easier
     // to use gemm operations from cublas.
     auto trainset =
-      make_device_mdarray<float>(handle, device_mr, make_extents<IdxT>(n_rows_train, dim));
+      make_device_mdarray<float>(handle, &managed_mr, make_extents<IdxT>(n_rows_train, dim));
 
     if constexpr (std::is_same_v<T, float>) {
       raft::spatial::knn::detail::utils::subsample(