fix a bug in batch-mode Coreset strategy (#32)

NUAA-AL · May 19, 2021 · 2ebf54c · 2ebf54c
1 parent 0679a56
commit 2ebf54c
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/alipy/query_strategy/query_labels.py b/alipy/query_strategy/query_labels.py
@@ -1837,6 +1837,8 @@ def update_distances(self, cluster_centers, only_new=True, reset_dist=False):
             if self.min_distances is None:
                 self.min_distances = np.min(dist, axis=1).reshape(-1, 1)
             else:
+                if dist.shape[1] != 1:
+                    dist = np.min(dist, axis=1).reshape(-1, 1)
                 self.min_distances = np.minimum(self.min_distances, dist)
 
     def select(self, label_index, unlabel_index, batch_size=1, **kwargs):
@@ -1879,6 +1881,7 @@ def select(self, label_index, unlabel_index, batch_size=1, **kwargs):
             assert self.train_idx[ind] in unlabel_index
 
             self.update_distances([ind], only_new=True, reset_dist=False)
+            self.min_distances[ind] = -100.0
             new_batch.append(self.train_idx[ind])
         return new_batch