Add missing device transfer in gpt_generate.py (#436)

This commit is contained in:
Sebastian Raschka 2024-11-14 19:12:53 +09:00 committed by GitHub
parent 6d1bb2cfb4
commit 129d0d740f

View File

@ -270,7 +270,7 @@ def main(gpt_config, input_prompt, model_size):
token_ids = generate(
model=gpt,
idx=text_to_token_ids(input_prompt, tokenizer),
idx=text_to_token_ids(input_prompt, tokenizer).to(device),
max_new_tokens=25,
context_size=gpt_config["context_length"],
top_k=50,