diff --git a/rayder/model.py b/rayder/model.py
index 6ff5eb3..c8dcc7c 100644
--- a/rayder/model.py
+++ b/rayder/model.py
@@ -359,12 +359,12 @@ def split(
             return t, torch.div(rem, cw, rounding_mode="floor"), rem.remainder(cw)
 
         def mask_mod(b, h, q_idx, kv_idx) -> Bool[torch.Tensor, ""]:
-            qt, qx, qy = split(q_idx)
-            kt_, kx, ky = split(kv_idx)
+            qt, qh, qw = split(q_idx)
+            kt_, kh_, kw_ = split(kv_idx)
             return (
                 ((qt.clamp(kt2, ct - 1 - kt2) - kt_).abs() <= kt2)
-                & ((qx.clamp(kw2, cw - 1 - kw2) - kx).abs() <= kw2)
-                & ((qy.clamp(kh2, ch - 1 - kh2) - ky).abs() <= kh2)
+                & ((qh.clamp(kh2, ch - 1 - kh2) - kh_).abs() <= kh2)
+                & ((qw.clamp(kw2, cw - 1 - kw2) - kw_).abs() <= kw2)
             )
 
         return mask_mod