randomize ties

LAION-AI · Mar 7, 2023 · e7173ec · e7173ec
1 parent af63b37
commit e7173ec
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/model/model_training/utils.py b/model/model_training/utils.py
@@ -93,7 +93,8 @@ def __iter__(self):
             epoch_idx.extend(sampled_idx)
 
         if self.samples_length is not None:
-            epoch_idx = sorted(epoch_idx, key=lambda x: -self.samples_length[x])
+            # sort by samples length and in case of ties randomize
+            epoch_idx = sorted(epoch_idx, key=lambda x: (self.samples_length[x], random.random()))
 
             if self.shuffle:
                 # do some minor shuffling to avoid repeating the same order