fix memory leak due to itertools.cycle

2024-04-05 10:59:32 +00:00
parent 5af00d0c1e
commit ad3379a73a
3 changed files with 28 additions and 2 deletions
--- a/lerobot/common/datasets/utils.py
+++ b/lerobot/common/datasets/utils.py
@@ -203,3 +203,12 @@ def compute_or_load_stats(dataset, batch_size=32, max_num_samples=None):

    torch.save(stats, stats_path)
    return stats
+
+
+def cycle(iterable):
+    iterator = iter(iterable)
+    while True:
+        try:
+            yield next(iterator)
+        except StopIteration:
+            iterator = iter(iterable)