From 7149ac30e60fdefee42541e2e6b47fc0f0011c48 Mon Sep 17 00:00:00 2001 From: Yao Matrix Date: Wed, 13 Mar 2024 17:09:53 +0800 Subject: [PATCH] Fix the issue of out of range (#98) Signed-off-by: yuanwu Co-authored-by: yuanwu --- server/text_generation_server/utils/tokens.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server/text_generation_server/utils/tokens.py b/server/text_generation_server/utils/tokens.py index 4eef456a..c4775a09 100644 --- a/server/text_generation_server/utils/tokens.py +++ b/server/text_generation_server/utils/tokens.py @@ -305,7 +305,7 @@ class HeterogeneousSampling: self.greedy = Greedy() def __call__(self, logits): - out = torch.empty(logits.shape[0], dtype=torch.int64, device=logits.device) + out = torch.zeros(logits.shape[0], dtype=torch.int64, device=logits.device) if self.greedy_indices: # Computing for all indices is faster than slicing torch.argmax(logits, -1, out=out)