Explicitly set return_dict for apply_chat_template (#33372)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
2026-01-30 07:27:04 +00:00
parent 726d89720c
commit 9432ed8c7e
11 changed files with 22 additions and 10 deletions
--- a/examples/offline_inference/prompt_embed_inference.py
+++ b/examples/offline_inference/prompt_embed_inference.py
@@ -38,8 +38,8 @@ def get_prompt_embeds(
    embedding_layer: torch.nn.Module,
 ):
    token_ids = tokenizer.apply_chat_template(
-        chat, add_generation_prompt=True, return_tensors="pt"
-    )
+        chat, add_generation_prompt=True, return_tensors="pt", return_dict=True
+    ).input_ids
    prompt_embeds = embedding_layer(token_ids).squeeze(0)
    return prompt_embeds