fixed bug in multi GPU setting

vllm-project · Mar 6, 2024 · e04e56d · e04e56d
1 parent c9a3db5
commit e04e56d
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/vllm/model_executor/models/jais.py b/vllm/model_executor/models/jais.py
@@ -93,8 +93,9 @@ def forward(
 
             # Get the logits for the next tokens.
             logits = self._get_logits(hidden_states, embedding, embedding_bias)
-            logits *= torch.tensor(float(output_logits_scale),
-                                   dtype=logits.dtype)
+            if logits is not None:
+                logits *= torch.tensor(float(output_logits_scale),
+                                       dtype=logits.dtype)
 
         # Only perform sampling in the driver worker.
         # Note: `_get_logits` is still distributed across TP workers because