Feat: Add multi-ec40-beam

speedcell4 · Sep 24, 2024 · 1a75c89 · 1a75c89
1 parent 1bd70ca
commit 1a75c89
Showing 1 changed file with 4 additions and 10 deletions.
diff --git a/torchglyph/data/loader.py b/torchglyph/data/loader.py
@@ -88,18 +88,12 @@ def new(cls, data_sources: Tuple[Dataset, ...],
         assert len(data_sources) > 0
 
         loaders = []
-        for index, (data_source, collate_fn, batch_size, sharding, drop_last, section_size, sortish_key) in enumerate(zip(
-                data_sources,
-                unpack(collate_fn),
-                unpack(batch_size),
-                unpack(sharding),
-                unpack(drop_last),
-                unpack(section_size),
-                unpack(sortish_key),
-        )):
+        for index, (data_source, collate_fn, batch_size, sharding, drop_last, section_size, sortish_key) in enumerate(
+                zip(data_sources, unpack(collate_fn), unpack(batch_size),
+                    unpack(sharding), unpack(drop_last), unpack(section_size), unpack(sortish_key))):
             training = index == 0
             if sharding:
-                data_source = data_source.select(range(get_rank(), len(data_source), get_world_size()))
+                data_source = data_source.shard(num_shards=get_world_size(), index=get_rank())
 
             sampler = SortishSampler(
                 data_source=data_source,