Add missing device transfer in gpt_generate.py

rasbt · Nov 14, 2024 · 95dd376 · 95dd376
1 parent 27a6a7e
commit 95dd376
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ch05/01_main-chapter-code/gpt_generate.py b/ch05/01_main-chapter-code/gpt_generate.py
@@ -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,