From ad93729dce5c0c730f2fbe6918e394ec71b75c59 Mon Sep 17 00:00:00 2001 From: Prince Canuma Date: Sun, 12 Jan 2025 15:23:19 +0100 Subject: [PATCH] Update llms/mlx_lm/models/cohere2.py Co-authored-by: Awni Hannun --- llms/mlx_lm/models/cohere2.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/llms/mlx_lm/models/cohere2.py b/llms/mlx_lm/models/cohere2.py index b1c9f2da..bc6ad276 100644 --- a/llms/mlx_lm/models/cohere2.py +++ b/llms/mlx_lm/models/cohere2.py @@ -157,7 +157,8 @@ class CohereModel(nn.Module): h = self.embed_tokens(inputs) if mask is None: - mask = create_attention_mask(h, cache, layer_idx=self.args.sliding_window_pattern - 1) + j = self.args.sliding_window_pattern + mask = create_attention_mask(h, cache[j - 1: j]) if cache is None: cache = [None] * len(self.layers)