Balanced sampler again

michaelbyun · Apr 18, 2023 · 9909466 · 9909466
1 parent bdbd7ce
commit 9909466
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/elk/extraction/prompt_loading.py b/elk/extraction/prompt_loading.py
@@ -178,7 +178,7 @@ def load_prompts(
  extra_cols = list(assert_type(Features, ds.features))
  extra_cols.remove(label_column)
 
- for example in BalancedSampler(ds, num_classes):
+ for example in BalancedSampler(ds, num_classes, label_col=label_column):
  yield _convert_to_prompts(
  example,
  label_column=label_column,