Skip to content

Commit f61c008

Browse files
authored
Add missing device transfer in gpt_generate.py (#436)
1 parent 27a6a7e commit f61c008

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

ch05/01_main-chapter-code/gpt_generate.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -270,7 +270,7 @@ def main(gpt_config, input_prompt, model_size):
270270

271271
token_ids = generate(
272272
model=gpt,
273-
idx=text_to_token_ids(input_prompt, tokenizer),
273+
idx=text_to_token_ids(input_prompt, tokenizer).to(device),
274274
max_new_tokens=25,
275275
context_size=gpt_config["context_length"],
276276
top_k=50,

0 commit comments

Comments
 (0)