BalancedBatchSampler reshuffle bugfix

adambielski · Mar 10, 2018 · 296164d · 296164d
1 parent 1f7c515
commit 296164d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/datasets.py b/datasets.py
@@ -178,7 +178,7 @@ def __iter__(self):
                                self.used_label_indices_count[class_]:self.used_label_indices_count[
                                                                          class_] + self.n_samples])
                 self.used_label_indices_count[class_] += self.n_samples
-                if self.used_label_indices_count[class_] + self.n_samples < len(self.label_to_indices[class_]):
+                if self.used_label_indices_count[class_] + self.n_samples > len(self.label_to_indices[class_]):
                     np.random.shuffle(self.label_to_indices[class_])
                     self.used_label_indices_count[class_] = 0
             yield indices