From 147d1297722d80bf8de3842cfba1863e041d23f4 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Sun, 25 Feb 2024 04:17:16 +1100 Subject: [PATCH] Update gemma.py --- unsloth/models/gemma.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index 8d40a63b66..b691588634 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -110,8 +110,8 @@ class FastGemmaRotaryEmbedding(torch.nn.Module): emb = torch.cat((freqs, freqs), dim=-1) seq_len = position_ids.shape[1] - new_cos = emb.cos().to(dtype=x.dtype)[:,:seq_len] - new_sin = emb.sin().to(dtype=x.dtype)[:,:seq_len] + new_cos = emb.cos().to(dtype=x.dtype)[:,:length] + new_sin = emb.sin().to(dtype=x.dtype)[:,:length] print(new_cos, new_cos.shape) print(old_cos, old_cos.shape)