bring back yield condition

YerevaNN · May 13, 2024 · 33f0567 · 33f0567
1 parent 254e2b7
commit 33f0567
Showing 1 changed file with 6 additions and 6 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