Update cpp/include/raft/neighbors/detail/cagra/cagra_search.cuh

Co-authored-by: Artem M. Chirkin <[email protected]>
rapidsai · Mar 13, 2024 · cf0db91 · cf0db91
1 parent 5d037b3
commit cf0db91
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/cpp/include/raft/neighbors/detail/cagra/cagra_search.cuh b/cpp/include/raft/neighbors/detail/cagra/cagra_search.cuh
@@ -266,10 +266,10 @@ void search_main(raft::resources const& res,
     // determined here. They are set just before kernel launch.
     using dataset_desc_t = standard_dataset_descriptor_t<T, InternalIdxT, 0, 0, DistanceT>;
     // Search using a plain (strided) row-major dataset
-    const dataset_desc_t dataset_desc(internal_dataset.data_handle(),
-                                      internal_dataset.extent(0),
-                                      internal_dataset.extent(1),
-                                      internal_dataset.stride(0));
+const dataset_desc_t dataset_desc(strided_dset->view().data_handle(),
+                                      strided_dset->n_rows(),
+                                      strided_dset->dim(),
+                                      strided_dset->stride());
 
     search_main_core<dataset_desc_t, CagraSampleFilterT>(
       res, params, dataset_desc, graph_internal, queries, neighbors, distances, sample_filter);