fix passing the wrong params for world size / rank

EleutherAI · Mar 28, 2023 · d2c66b0 · d2c66b0
1 parent 764fda9
commit d2c66b0
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)