reword conditional

rapidsai · Nov 12, 2024 · 11a7402 · 11a7402
1 parent 14b14e8
commit 11a7402
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/cpp/src/neighbors/detail/nn_descent.cuh b/cpp/src/neighbors/detail/nn_descent.cuh
@@ -473,7 +473,8 @@ RAFT_KERNEL preprocess_data_kernel(
   if (threadIdx.x == 0) { l2_norm = 0; }
   __syncthreads();
 
-  if (metric != cuvs::distance::DistanceType::InnerProduct) {
+  if (metric == cuvs::distance::DistanceType::L2Expanded ||
+      metric == cuvs::distance::DistanceType::CosineExpanded) {
     int lane_id = threadIdx.x % raft::warp_size();
     for (int step = 0; step < raft::ceildiv(dim, raft::warp_size()); step++) {
       int idx         = step * raft::warp_size() + lane_id;