Fix gather_if raft update (#5149)

Authors: - Micka (https://github.com/lowener) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: #5149
rapidsai · Jan 23, 2023 · de32125 · de32125
1 parent 272f856
commit de32125
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/cpp/src/kmeans/kmeans_mg_impl.cuh b/cpp/src/kmeans/kmeans_mg_impl.cuh
@@ -636,7 +636,7 @@ void fit(const raft::handle_t& handle,
       centroids.extent(0),
       itr_wt,
       itr_wt,
-      wtInCluster.size(),
+      wtInCluster.extent(0),
       newCentroids.data_handle(),
       [=] __device__(raft::KeyValuePair<ptrdiff_t, DataT> map) {  // predicate
         // copy when the # of samples in the cluster is 0