From 52afdcc281187f2cb1602e553d70a72c1874fa45 Mon Sep 17 00:00:00 2001 From: Cyril Vallez Date: Mon, 20 Jan 2025 15:25:10 +0100 Subject: [PATCH] update comment --- .../models/transformers_flash_causal_lm.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/server/text_generation_server/models/transformers_flash_causal_lm.py b/server/text_generation_server/models/transformers_flash_causal_lm.py index 30ea4c8f..17f47e5e 100644 --- a/server/text_generation_server/models/transformers_flash_causal_lm.py +++ b/server/text_generation_server/models/transformers_flash_causal_lm.py @@ -243,10 +243,8 @@ class TransformersFlashCausalLM(FlashCausalLM): adapter_data=None, # not supported, but passed to match original signature ): hidden_states = self.model.model.forward( - input_ids=input_ids.unsqueeze(0), # expand dim to easily fit transformers - position_ids=position_ids.unsqueeze( - 0 - ), # expand dim to easily fit transformers + input_ids=input_ids.unsqueeze(0), # expand dim to fit Transformers + position_ids=position_ids.unsqueeze(0), # expand dim to fit Transformers past_key_values=None, # we use self.kv_cache instead of transformers cache object use_cache=False, # we use self.kv_cache instead of transformers cache object return_dict=True,