Fix dataset dimension in IVF-PQ C wrappers (#292)

fixes #288 Authors: - Tamas Bela Feher (https://github.com/tfeher) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: #292
rapidsai · Aug 14, 2024 · 9cf4800 · 9cf4800
1 parent 829ae16
commit 9cf4800
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/cpp/src/neighbors/ivf_flat_c.cpp b/cpp/src/neighbors/ivf_flat_c.cpp
@@ -47,7 +47,7 @@ void* _build(cuvsResources_t res, cuvsIvfFlatIndexParams params, DLManagedTensor
   build_params.conservative_memory_allocation = params.conservative_memory_allocation;
 
   auto dataset = dataset_tensor->dl_tensor;
-  auto dim     = dataset.shape[0];
+  auto dim     = dataset.shape[1];
 
   auto index = new cuvs::neighbors::ivf_flat::index<T, IdxT>(*res_ptr, build_params, dim);
 

diff --git a/cpp/src/neighbors/ivf_pq_c.cpp b/cpp/src/neighbors/ivf_pq_c.cpp
@@ -51,7 +51,7 @@ void* _build(cuvsResources_t res, cuvsIvfPqIndexParams params, DLManagedTensor*
   build_params.max_train_points_per_pq_code   = params.max_train_points_per_pq_code;
 
   auto dataset = dataset_tensor->dl_tensor;
-  auto dim     = dataset.shape[0];
+  auto dim     = dataset.shape[1];
 
   auto index = new cuvs::neighbors::ivf_pq::index<IdxT>(*res_ptr, build_params, dim);