fix: fixed online loader sharding across processes

AshishKumar4 · Sep 11, 2024 · 9f4eee0 · 9f4eee0
1 parent 148edb8
commit 9f4eee0
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/flaxdiff/data/online_loader.py b/flaxdiff/data/online_loader.py
@@ -214,8 +214,9 @@ def parallel_image_loader(
         iteration = 0
         while True:
             # Repeat forever
-            shards = [dataset[i*shard_len:(i+1)*shard_len]
-                      for i in range(num_workers)]
+            # shards = [dataset[i*shard_len:(i+1)*shard_len]
+            #           for i in range(num_workers)]
+            shards = [dataset.shard(num_shards=num_workers, index=i) for i in range(num_workers)]
             print(f"mapping {len(shards)} shards")
             pool.map(map_batch_fn, shards)
             iteration += 1

diff --git a/setup.py b/setup.py
@@ -11,7 +11,7 @@
 setup(
     name='flaxdiff',
     packages=find_packages(),
-    version='0.1.34',
+    version='0.1.35',
     description='A versatile and easy to understand Diffusion library',
     long_description=open('README.md').read(),
     long_description_content_type='text/markdown',