Fix VPQ search params validation

rapidsai · Mar 20, 2024 · 1fb7c36 · achirkin · Mar 20, 2024 · 1fb7c36
1 parent 103b9c0
commit 1fb7c36
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/cpp/include/raft/neighbors/detail/cagra/cagra_search.cuh b/cpp/include/raft/neighbors/detail/cagra/cagra_search.cuh
@@ -165,6 +165,9 @@ void launch_vpq_search_main_core(
   raft::device_matrix_view<DistanceT, int64_t, row_major> distances,
   CagraSampleFilterT sample_filter)
 {
+  assert(vpq_dset->pq_bits() == 2 || vpq_dset->pq_bits() == 4);
+  assert(vpq_dset->dim() % vpq_dset->pq_dim() == 0);
+
   const float vq_scale = 1.0f;
   const float pq_scale = 1.0f;