Fix auto deploy demo logprobs return issue.

Signed-off-by: Wangshanshan <[email protected]>

Author: dominicshanshan

tensorrt_llm/_torch/auto_deploy/shim/demollm.py

@@ -235,11 +235,11 @@ def _sample(
         logits_shape = logits.shape
         logits = logits.view(-1, logits_shape[-1])  # sampling_batch expects 2D logits
         if isinstance(sampling_params.top_k, int) and sampling_params.top_k > 1:
-            idx_next, probs = top_k_sampling_batch(
+            idx_next, probs, _ = top_k_sampling_batch(
                 logits, top_k=sampling_params.top_k, temperature=1.0
             )
         else:
-            idx_next, probs = greedy_search_sampling_batch(logits)
+            idx_next, probs, _ = greedy_search_sampling_batch(logits)
         idx_next = idx_next.view(logits_shape[:-1])
         return idx_next, probs