From 0dce4ffc2da057b22e0d83487909be8b25aaa628 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 1 Feb 2025 19:15:55 -0800 Subject: [PATCH] Update gemma.py --- unsloth/models/gemma.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index d94f24071b..23561ed07e 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -211,12 +211,11 @@ class GemmaFixedRotaryEmbedding(torch.nn.Module): ): super().__init__() if config is not None: - print(config) - print(dir(config)) # [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)) + dim = getattr(config, "head_dim", None) + if dim is None: dim = int((config.hidden_size // config.num_attention_heads)) device = "cuda" max_position_embeddings = config.max_position_embeddings pass