From caf252f193affc39e47482f80f49311345ce4440 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 1 Feb 2025 17:54:00 -0800 Subject: [PATCH] Update gemma.py --- unsloth/models/gemma.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index c654343282..408c55440c 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -210,7 +210,14 @@ class GemmaFixedRotaryEmbedding(torch.nn.Module): config = None, # [TODO] Hack to pass in config - need to remove later ): super().__init__() - if config is not None: return # [TODO] Hack to pass in config - need to remove later + if config is not None: + # [TODO] Hack to pass in config - need to remove later + base = config.rope_theta + partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0 + dim = int((config.hidden_size // config.num_attention_heads)) + device = "cuda" + max_position_embeddings = config.max_position_embeddings + pass self.dim = dim self.max_position_embeddings = max_position_embeddings self.base = base