Update vllm/model_executor/models/granitemoe.py

Co-authored-by: Nick Hill <[email protected]>
vllm-project · Sep 30, 2024 · 2d86b7d · 2d86b7d
1 parent 280b22b
commit 2d86b7d
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/vllm/model_executor/models/granitemoe.py b/vllm/model_executor/models/granitemoe.py
@@ -382,10 +382,11 @@ def forward(
         return hidden_states
 
     def compute_logits(self, hidden_states: torch.Tensor,
-                       sampling_metadata: SamplingMetadata) -> torch.Tensor:
+                       sampling_metadata: SamplingMetadata) -> Optional[torch.Tensor]:
         logits = self.logits_processor(self.lm_head, hidden_states,
                                        sampling_metadata)
-        logits = logits / self.config.logits_scaling
+        if logits is not None:
+            logits /= self.config.logits_scaling
         return logits
 
     def make_empty_intermediate_tensors(