Fix Kmeans cluster updates issue

src/torch_kmeans/clustering/constr_kmeans.py

@@ -312,7 +312,7 @@ def _cluster(
             # get cluster assignments
             c_assign = self._assign(x, centers, weights, k_mask)
             # update cluster centers
-            centers = group_by_label_mean(x, c_assign, k_max_range)
+            centers = group_by_label_mean(x, c_assign, old_centers, k_max_range)
             if self.tol is not None:
                 # calculate center shift
                 shift = self._calculate_shift(centers, old_centers, p=self.p_norm)

src/torch_kmeans/clustering/kmeans.py

@@ -541,7 +541,7 @@ def _cluster(
             # get cluster assignments
             c_assign = self._assign(x, centers)
             # update cluster centers
-            centers = group_by_label_mean(x, c_assign, k_max_range)
+            centers = group_by_label_mean(x, c_assign, old_centers, k_max_range)
             if self.tol is not None:
                 # calculate center shift
                 shift = self._calculate_shift(centers, old_centers, p=self.p_norm)

src/torch_kmeans/utils/utils.py

@@ -34,6 +34,7 @@ class ClusterResult(NamedTuple):
 def group_by_label_mean(
     x: Tensor,
     labels: Tensor,
+    old_centers: Tensor,
     k_max_range: Tensor,
 ) -> Tensor:
     """Group samples in x by label
@@ -64,7 +65,11 @@ def group_by_label_mean(
         .to(x.dtype)
     )
     M = F.normalize(M, p=1.0, dim=-1)
-    return torch.matmul(M, x[:, None, :, :].expand(bs, m, n, d))
+    new_centers = torch.matmul(M, x[:, None, :, :].expand(bs, m, n, d))
+    nan_mask = torch.isnan(new_centers)
+    new_centers = torch.where(nan_mask, old_centers, new_centers)
+    # return torch.matmul(M, x[:, None, :, :].expand(bs, m, n, d))
+    return new_centers
 
 
 @torch.jit.script