Formatting & adding quant_config

vllm-project · Aug 19, 2024 · 553e89c · 553e89c
1 parent a5c7b7a
commit 553e89c
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/vllm/model_executor/models/granite.py b/vllm/model_executor/models/granite.py
@@ -278,6 +278,7 @@ def __init__(
                 self.vocab_size,
                 config.hidden_size,
                 org_num_embeddings=config.vocab_size,
+                quant_config=quant_config,
             )
         else:
             self.embed_tokens = PPMissingLayer()
@@ -424,8 +425,9 @@ def forward(
                                   attn_metadata, intermediate_tensors)
         return model_output
 
-    def compute_logits(self, hidden_states: torch.Tensor,
-                       sampling_metadata: SamplingMetadata) -> Optional[torch.Tensor]:
+    def compute_logits(
+            self, hidden_states: torch.Tensor,
+            sampling_metadata: SamplingMetadata) -> Optional[torch.Tensor]:
         logits = self.logits_processor(self.lm_head, hidden_states,
                                        sampling_metadata)
         logits /= self.config.logits_scaling