diff --git a/llms/mlx_lm/models/cohere2.py b/llms/mlx_lm/models/cohere2.py
index 784bcfcf..d489d2b9 100644
--- a/llms/mlx_lm/models/cohere2.py
+++ b/llms/mlx_lm/models/cohere2.py
@@ -156,13 +156,13 @@ class CohereModel(nn.Module):
     ):
         h = self.embed_tokens(inputs)
 
+        if cache is None:
+            cache = [None] * len(self.layers)
+            
         if mask is None:
             j = self.args.sliding_window_pattern
             mask = create_attention_mask(h, cache[j - 1 : j])
 
-        if cache is None:
-            cache = [None] * len(self.layers)
-
         for layer, c in zip(self.layers, cache):
             h = layer(h, mask, c)