Fix bug of cycling nums

yzhangcs · May 31, 2023 · 5833ee8 · 5833ee8
1 parent d35fd45
commit 5833ee8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/supar/utils/data.py b/supar/utils/data.py
@@ -264,7 +264,7 @@ def cycle(length):
                 for i in range_fn(length).tolist():
                     yield i
 
-        for i in cycle(range(len(self.buckets))):
+        for i in cycle(len(self.buckets)):
             bucket = self.buckets[i]
             split_sizes = [(len(bucket) - j - 1) // self.n_batches[i] + 1 for j in range(self.n_batches[i])]
             # DON'T use `torch.chunk` which may return wrong number of batches