Fix k-means++ mini-batch initialization (#22)

borchero · Jun 17, 2022 · 1f73128 · 1f73128
1 parent 1eb6af0
commit 1f73128
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pycave/clustering/kmeans/estimator.py b/pycave/clustering/kmeans/estimator.py
@@ -112,7 +112,7 @@ def fit(self, data: TensorLike) -> KMeans:
             batch_size=self.batch_size or len(data),
             collate_fn=collate_tensor,
         )
-        is_batch_training = self._num_batches_per_epoch(loader) == 1
+        is_batch_training = self._num_batches_per_epoch(loader) > 1
 
         # First, initialize the centroids
         if self.init_strategy == "random":