ndim check in llama

cb575f5dc34a12397279f3a7e1b138a68366a8cf · 2024-09-23 12:31:25 +0100 · Alex Cheema

Files touched

M exo/inference/mlx/models/llama.py

Diff

commit cb575f5dc34a12397279f3a7e1b138a68366a8cf
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Mon Sep 23 12:31:25 2024 +0100

    ndim check in llama
---
 exo/inference/mlx/models/llama.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/exo/inference/mlx/models/llama.py b/exo/inference/mlx/models/llama.py
index 719d6a88..11b6716a 100644
--- a/exo/inference/mlx/models/llama.py
+++ b/exo/inference/mlx/models/llama.py
@@ -54,7 +54,7 @@ class LlamaModel(nn.Module):
       h = inputs
 
     mask = None
-    if h.shape[1] > 1:
+    if h.ndim > 1 and h.shape[1] > 1:
       mask = create_attention_mask(h, cache)
 
     if cache is None:

← 7dd7fe49 fix allow patterns · back to Exo · clarify which models exo supports 2244ff4a →