From 7cb49f6f4f47d13bfe022fbc14a977e3eb329f9d Mon Sep 17 00:00:00 2001 From: Mohit Sharma Date: Fri, 27 Sep 2024 15:53:44 +0000 Subject: [PATCH] float16 dep --- .../models/custom_modeling/flash_mistral_modeling.py | 1 + 1 file changed, 1 insertion(+) 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 index 3e16d371..341a2352 100644 --- 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 @@ class MistralMLP(nn.Module): if ( SYSTEM == "rocm" and self.hidden_act == "silu" + and hidden_states.dtype == torch.float16 and hidden_states.shape[0] == 1 and not self.quantize ):