From 1a8a18d541b7557be2c0da3994b77d8d7c17b0f0 Mon Sep 17 00:00:00 2001 From: Nicolas Patry Date: Wed, 8 May 2024 06:33:13 +0000 Subject: [PATCH] Cleanup. --- server/text_generation_server/utils/layers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server/text_generation_server/utils/layers.py b/server/text_generation_server/utils/layers.py index 9f6f6c1e..af00f5a3 100644 --- a/server/text_generation_server/utils/layers.py +++ b/server/text_generation_server/utils/layers.py @@ -545,7 +545,7 @@ class MLPSpeculatorModel(torch.nn.Module): # h indicates # of generated tokens state = hidden_states b = state.size(0) - ind = input_ids[-b:].unsqueeze(0) + ind = input_ids.unsqueeze(0) out = torch.empty(1, b, self.n_predict, device=state.device).int() # b k h # log_probs = torch.zeros(1, b, device=state.device) # b k all_probs = torch.empty(