diff --git a/elk/extraction/extraction.py b/elk/extraction/extraction.py
index 5446cd5c..17a42ffa 100644
--- a/elk/extraction/extraction.py
+++ b/elk/extraction/extraction.py
@@ -191,6 +191,8 @@ def extract_hiddens(
     prompt_ds = load_prompts(
         ds_names[0],
         binarize=cfg.binarize,
+        num_shots=cfg.num_shots,
+        num_variants=cfg.num_variants,
         split_type=split_type,
         template_path=cfg.template_path,
         rank=rank,