Feat: Update RandomSortishSampler reserve

speedcell4 · Nov 28, 2024 · afa10e3 · afa10e3
1 parent ec9fc60
commit afa10e3
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/torchglyph/data/sampler.py b/torchglyph/data/sampler.py
@@ -39,7 +39,12 @@ def __iter__(self):
 
 class RandomSortishSampler(_SortishSampler):
     def __iter__(self):
-        idx, key, reverse = [], [], True
+        idx, key = [], []
+
+        if distributed.is_initialized():
+            reverse = distributed.get_rank() % 2 == 0
+        else:
+            reverse = True
 
         while True:
             for batch in self.ds.shuffle().iter(batch_size=self.section_size, drop_last_batch=False):