Moved the encoded_prompts to correct device

This commit is contained in:
alberduris
2019-12-31 13:11:41 +01:00
committed by Lysandre Debut
parent 1efc208ff3
commit dd4df80f0b
285 changed files with 1 additions and 0 deletions

1
examples/run_generation.py Normal file → Executable file
View File

@@ -212,6 +212,7 @@ def main():
prepare_input = PREPROCESSING_FUNCTIONS.get(args.model_type)
prompt_text = prepare_input(args, model, tokenizer, prompt_text)
encoded_prompt = tokenizer.encode(prompt_text, add_special_tokens=False, return_tensors="pt")
encoded_prompt = encoded_prompt.to(args.device)
output_sequences = model.generate(
input_ids=encoded_prompt,