From 3695be30b5544978ef64d6cc9398df3b568fbe37 Mon Sep 17 00:00:00 2001 From: namgyaaal Date: Fri, 3 Oct 2025 11:34:20 -0500 Subject: [PATCH] Fix bug in MultiHeadAttention --- parakeet_mlx/attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/parakeet_mlx/attention.py b/parakeet_mlx/attention.py index 7cfdc3c..1f611a6 100644 --- a/parakeet_mlx/attention.py +++ b/parakeet_mlx/attention.py @@ -44,7 +44,7 @@ def __call__( k, v = cache.update_and_fetch_kv(k, v) o = mx.fast.scaled_dot_product_attention(q, k, v, scale=self.scale, mask=mask) - o = o.transpose(0, 2, 1, 3).reshape(batch, q_seq, self.n_feat) + o = o.transpose(0, 2, 1, 3).reshape(batch, q_seq, self.head_dim * self.n_head) return self.linear_out(o)