This commit is contained in:
Prince Canuma 2025-03-12 10:30:02 +01:00
parent 2d30f6787a
commit 37cd4a382b

View File

@ -100,7 +100,7 @@ class Attention(nn.Module):
if self.is_sliding and mask is not None: if self.is_sliding and mask is not None:
key_len = keys.shape[-2] key_len = keys.shape[-2]
if mask.shape[-1] != key_len: if mask.shape[-1] != key_len:
mask = mask[..., :key_len] mask = mask[..., -key_len:]
output = mx.fast.scaled_dot_product_attention( output = mx.fast.scaled_dot_product_attention(