Merge branch 'multi-datasets' into james-debug

EleutherAI · Mar 28, 2023 · 0c191e3 · 0c191e3
2 parents 272e8e8 + d2c66b0
commit 0c191e3
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/elk/extraction/prompt_loading.py b/elk/extraction/prompt_loading.py
@@ -130,7 +130,9 @@ def load_prompts(
         # only keep the datapoints relevant to the current process
         if world_size > 1:
             # This prints to stdout which is slightly annoying
-            split = split_dataset_by_node(split, world_size, rank)
+            split = split_dataset_by_node(
+                dataset=split, rank=rank, world_size=world_size
+            )
 
         raw_datasets.append(split)
         train_datasets.append(train_ds)