float16 dep

huggingface · Sep 27, 2024 · 7cb49f6 · 7cb49f6
1 parent b2cd1b6
commit 7cb49f6
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/server/text_generation_server/models/custom_modeling/flash_mistral_modeling.py b/server/text_generation_server/models/custom_modeling/flash_mistral_modeling.py
@@ -303,6 +303,7 @@ def forward(self, hidden_states, adapter_data):
         if (
             SYSTEM == "rocm"
             and self.hidden_act == "silu"
+            and hidden_states.dtype == torch.float16
             and hidden_states.shape[0] == 1
             and not self.quantize
         ):