fix: small bug in _updated_centroids() in kmeans.py, wrong initializa…

…tion of zero_counts (#1885) Spotted a one-line bug in _updated_centroids() Co-authored-by: Lei Xu <[email protected]>
lancedb · Jan 31, 2024 · 1a67a17 · 1a67a17
1 parent 5807cde
commit 1a67a17
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/python/lance/torch/kmeans.py b/python/python/lance/torch/kmeans.py
@@ -174,7 +174,7 @@ def _updated_centroids(
         self, centroids: torch.Tensor, counts: torch.Tensor
     ) -> torch.Tensor:
         centroids = centroids / counts[:, None]
-        zero_counts = centroids == 0
+        zero_counts = counts == 0
         for idx in zero_counts.nonzero(as_tuple=False):
             # split the largest cluster and remove empty cluster
             max_idx = torch.argmax(counts).item()