update func export_codec_vllm

2026-02-05 09:59:23 +08:00 · 2025-02-26 16:48:21 +08:00
parent f6a18ee07a
commit f280558bcb
3 changed files with 5 additions and 2 deletions
--- a/cosyvoice/llm/llm.py
+++ b/cosyvoice/llm/llm.py
@@ -343,7 +343,7 @@ class Qwen2LM(TransformerLM):
                                            max_tokens=max_len)
            request_id = uuid.uuid4()
            vllm_codec_engine.add_request(request_id,
-                                        {"prompt_embeds": lm_input.to(torch.bfloat16).to(device)},
+                                        {"prompt_embeds": lm_input.squeeze(0).to(torch.bfloat16).to(device)},
                                        sampling_params)
            ## generator
            out_token_ids = []