diff --git a/chemlactica/jsonl_dataset.py b/chemlactica/jsonl_dataset.py
index 8c79423..9fd0cd8 100644
--- 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
index fe49900..8ead52f 100644
--- a/test_status.yaml
+++ b/test_status.yaml
@@ -1 +1 @@
-16b42bfd74750347376025e520905710f7b791d8: PASS
+33f0567f336bd041b6b687cb3258a855d948b6b8: PASS