Fix logit indices

2026-03-29 02:51:27 +08:00 · 2024-04-17 18:01:43 +00:00 · 2024-04-17 18:01:43 +00:00 · 8888d1c474
commit 8888d1c474
parent cedb67028a
1 changed files with 1 additions and 1 deletions
--- a/vllm/model_executor/models/jax/gemma.py
+++ b/vllm/model_executor/models/jax/gemma.py
@ -311,7 +311,7 @@ class Transformer(nn.Module):
      )
      kv_caches[i] = layer_cache
    x = self.final_norm(x)
-
+    x = x.reshape(-1, x.shape[-1])
    hidden_states = x[logits_indices]
    logits = self.embedder.decode(hidden_states)
    return logits, kv_caches