From fead2e2146c224336c7f31851a3f7784325dd818 Mon Sep 17 00:00:00 2001 From: Tom Stesco Date: Thu, 31 Oct 2024 03:44:46 +0000 Subject: [PATCH] fixing #31 by converting SamplerOutput output_token to int in TTModelRunner::_make_sampler_output as expected by vLLM downstream --- vllm/worker/tt_model_runner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/worker/tt_model_runner.py b/vllm/worker/tt_model_runner.py index 3a8478bfd952d..0018e4d42ef59 100644 --- a/vllm/worker/tt_model_runner.py +++ b/vllm/worker/tt_model_runner.py @@ -351,7 +351,7 @@ def _make_sampler_output( zero_logprob = Logprob(0.0) sampler_outputs = [] for batch_idx, seq_id in enumerate(seq_groups): - next_token_id = next_token_ids[batch_idx] + next_token_id = int(next_token_ids[batch_idx]) seq_outputs = [SequenceOutput(seq_id, next_token_id, {next_token_id: zero_logprob})] sampler_outputs.append(