Fix mistral after syn1.15 update (#858)

huggingface · Apr 3, 2024 · eaac913 · eaac913
1 parent 445be21
commit eaac913
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/optimum/habana/transformers/models/mistral/modeling_mistral.py b/optimum/habana/transformers/models/mistral/modeling_mistral.py
@@ -518,7 +518,7 @@ def forward(
 
 
 class GaudiMistralForCausalLM(MistralForCausalLM):
-    def allocate_kv_cache(self, batch_size, seq_len, _, __):
+    def allocate_kv_cache(self, batch_size, seq_len, _):
         self.model.allocate_kv_cache(batch_size, seq_len)
 
     def reorder_kv_cache(self, beam_idx: torch.LongTensor):