From de321257685012a49a65f9112875481c1e3e15a4 Mon Sep 17 00:00:00 2001
From: Micka <9810050+lowener@users.noreply.github.com>
Date: Tue, 24 Jan 2023 00:11:48 +0100
Subject: [PATCH] Fix gather_if raft update (#5149)

Authors:
  - Micka (https://github.com/lowener)

Approvers:
  - Corey J. Nolet (https://github.com/cjnolet)

URL: https://github.com/rapidsai/cuml/pull/5149
---
 cpp/src/kmeans/kmeans_mg_impl.cuh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cpp/src/kmeans/kmeans_mg_impl.cuh b/cpp/src/kmeans/kmeans_mg_impl.cuh
index 0fdfc6b826..c17cc9467b 100644
--- 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