From f1041f885f7c0bc214c7c457ebbafbe94a2beeec Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Th=C3=A9ophile=20Lafargue?= <138336683+eauchs@users.noreply.github.com> Date: Thu, 9 Apr 2026 14:56:42 +0200 Subject: [PATCH] fix(gemma4): cast RoPE offset to int before mx.arange() (#4901) * fix(gemma4): cast RoPE offset to int before mx.arange() * fix(gemma4): use zero-based arange + offset to avoid CPU-GPU sync --- unsloth/models/gemma4_text.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/unsloth/models/gemma4_text.py b/unsloth/models/gemma4_text.py index af8bb29346..fd6af0307b 100644 --- a/unsloth/models/gemma4_text.py +++ b/unsloth/models/gemma4_text.py @@ -247,7 +247,7 @@ class ProportionalRoPE(nn.Module): def __call__(self, x: mx.array, offset: int = 0) -> mx.array: # x shape: (B, n_heads, L, head_dim) seq_len = x.shape[-2] - positions = mx.arange(offset, offset + seq_len, dtype = mx.float32) + positions = mx.arange(seq_len, dtype = mx.float32) + offset # (L, head_dim//2) freqs = mx.outer(positions, self._inv_freq)