[Speculative decoding] [Bugfix] Fix overallocation in ngram + spec logprobs (#4672)

This commit is contained in:
Cade Daniel
2024-05-08 16:24:46 -07:00
committed by GitHub
parent f942efb5a3
commit 8b9241be3a

View File

@ -138,7 +138,7 @@ class NGramWorker(LoraNotSupportedWorkerBase):
SamplerOutput(
outputs=None,
sampled_token_probs=token_probs[i],
logprobs=token_logprobs,
logprobs=token_logprobs[i],
sampled_token_ids=token_ids[i],
))
return outputs, False