Fix unique (PaddlePaddle#60840)

* fix unique kernel, row to num_out
warrentdrew · Jan 23, 2024 · 0bc8d81 · 0bc8d81
1 parent 1aa5f4b
commit 0bc8d81
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/paddle/phi/kernels/gpu/unique_kernel.cu b/paddle/phi/kernels/gpu/unique_kernel.cu
@@ -388,9 +388,11 @@ static void ComputeUniqueDims(const Context& context,
   // 3. counts: 'counts'
   counts->Resize(common::make_ddim({num_out}));
   auto* count_data = context.template Alloc<IndexT>(counts);
-  thrust::fill(exec_policy, count_data, count_data + row, 0);
-  thrust::adjacent_difference(
-      exec_policy, range_data_ptr + 1, range_data_ptr + row + 1, count_data);
+  thrust::fill(exec_policy, count_data, count_data + num_out, 0);
+  thrust::adjacent_difference(exec_policy,
+                              range_data_ptr + 1,
+                              range_data_ptr + num_out + 1,
+                              count_data);
 }
 
 // Calculate unique when 'axis' is set