From 1f7312848608765ce387a40211ae62c41e1f077c Mon Sep 17 00:00:00 2001 From: Oliver Borchert Date: Fri, 17 Jun 2022 18:55:54 +0200 Subject: [PATCH] Fix k-means++ mini-batch initialization (#22) --- pycave/clustering/kmeans/estimator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pycave/clustering/kmeans/estimator.py b/pycave/clustering/kmeans/estimator.py index 08646eb..df6ed99 100644 --- 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":