Use managed memory for the new temprorary buffer

rapidsai · Jan 18, 2024 · f322d0a · f322d0a
1 parent 0cb6caf
commit f322d0a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/cpp/include/raft/neighbors/detail/ivf_pq_build.cuh b/cpp/include/raft/neighbors/detail/ivf_pq_build.cuh
@@ -1772,7 +1772,7 @@ auto build(raft::resources const& handle,
       // TODO(tfeher): Enable codebook generation with any type T, and then remove
       // trainset tmp.
       auto trainset_tmp =
-        make_device_mdarray<T>(handle, device_mr, make_extents<IdxT>(n_rows_train, dim));
+        make_device_mdarray<T>(handle, &managed_mr, make_extents<IdxT>(n_rows_train, dim));
       raft::spatial::knn::detail::utils::subsample(
         handle, dataset, n_rows, trainset_tmp.view(), random_seed);
       cudaDeviceSynchronize();