Merge pull request #27 from YerevaNN/remove_dload_dup

Remove dataloader duplication via yield on each process
YerevaNN · May 14, 2024 · 7be6d91 · 7be6d91
2 parents 254e2b7 + ea2c700
commit 7be6d91
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/chemlactica/jsonl_dataset.py b/chemlactica/jsonl_dataset.py
@@ -50,12 +50,12 @@ def samples_generator(
                 counter = 0
                 while line:
                     state["position"] = f.tell()
-                    # if should_yield_on_current_rank(
-                    #     counter,
-                    #     distributed_state.num_processes,
-                    #     distributed_state.process_index,
-                    # ):
-                    #     returned = True
+                    if should_yield_on_current_rank(
+                        counter,
+                        distributed_state.num_processes,
+                        distributed_state.process_index,
+                    ):
+                        returned = True
                     ret = format_sample(line)
                     yield ret
                     counter = counter + 1

diff --git a/test_status.yaml b/test_status.yaml
@@ -1 +1 @@
-16b42bfd74750347376025e520905710f7b791d8: PASS
+33f0567f336bd041b6b687cb3258a855d948b6b8: PASS