hotfix: fix flashllama

huggingface · Oct 23, 2024 · 27ff187 · 27ff187
1 parent 03c9388
commit 27ff187
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/server/text_generation_server/models/custom_modeling/flash_llama_modeling.py b/server/text_generation_server/models/custom_modeling/flash_llama_modeling.py
@@ -692,7 +692,7 @@ def forward(
         logits, speculative_logits = self.lm_head(hidden_states)
 
         # Used in Granite
-        if not self.logits_scaled:
+        if self.logits_scaling is not None and not self.logits_scaled:
             logits /= self.logits_scaling
             if speculative_logits is not None:
                 speculative_logits /= self.logits_scaling