From 3137392ccf048a7316bbf1b2204f49ace393c540 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Tue, 5 Mar 2024 02:41:56 +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 b32357a5bb..c99bd5f927 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -289,9 +289,9 @@ class GemmaFixedRotaryEmbedding(torch.nn.Module): timescale = self.base**freq_exponents positions = torch.arange(self.max_seq_len_cached, device = "cpu", dtype = torch.int64).float() radians_new = positions[..., None] / timescale[None, None, :] - + radians_new = radians_new.squeeze(0) + emb = torch.cat((radians_new, radians_new), dim=-1) - print(emb.shape) self.register_buffer("cos_cached", emb.cos().to(dtype=dtype, device=device, non_blocking=True), persistent=False) self.register_buffer("sin_cached", emb.sin().to(dtype=dtype, device=device, non_blocking=True), persistent=False) pass