Skip to content

Commit 94a7778

Browse files
authored
explicitly set the task in the HF pipeline so it can work in offline mode (#31)
1 parent 8430abe commit 94a7778

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

generate/generate.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -98,7 +98,7 @@
9898
prompts_repeated = [p for p in prompts for _ in range(args.num_samples_per_prompt)]
9999

100100
""" Initialize HuggingFace pipeline for generation """
101-
generator = pipeline(model=args.model, torch_dtype=inference_config.get_dtype(), device=0, token=args.hf_token)
101+
generator = pipeline(task="text-generation", model=args.model, torch_dtype=inference_config.get_dtype(), device=0, token=args.hf_token)
102102
inference_config.init_padding(generator.tokenizer)
103103

104104
""" Create a prompt data set to pass to generate method """

0 commit comments

Comments
 (0)