Merge branch 'habana-main' into 2.3.0

huggingface · Oct 27, 2024 · c23584f · c23584f
2 parents 372e071 + b126bf4
commit c23584f
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/server/text_generation_server/models/causal_lm.py b/server/text_generation_server/models/causal_lm.py
@@ -731,13 +731,11 @@ def __init__(
         }
 
 
-        if model.config.model_type in ["llama", "mistral", "starcoder2", "qwen2", "falcon", "gemma"]:
-
+        if model.config.model_type in ["llama", "mistral", "starcoder2", "qwen2", "falcon"]:
             if model.config.model_type not in ["falcon"]:
                 self.kwargs["attn_softmax_bf16"] = True
 
-            if model.config.model_type not in ["gemma"]:
-                self.kwargs["trim_logits"] = True
+            self.kwargs["trim_logits"] = True
 
             if os.getenv("USE_FLASH_ATTENTION", "false").lower() == "true":
                 self.kwargs["use_flash_attention"] = True