diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py
index 23561ed07e..bc29c46abc 100644
--- a/unsloth/models/gemma.py
+++ b/unsloth/models/gemma.py
@@ -224,7 +224,6 @@ class GemmaFixedRotaryEmbedding(torch.nn.Module):
         self.base = base
         # Dynamic RoPE we first set it to a max of 4 * 8192 tokens then we iteratively grow this
         self.current_rope_size = min(4 * 8192, self.max_position_embeddings)
-        print(dim, max_position_embeddings, base)
 
         # Build here to make `torch.jit.trace` work.
         self._set_cos_sin_cache(seq_len=self.current_rope_size, device=device, dtype=torch.get_default_dtype())
diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index da3295adfd..4b64c74f3e 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -1159,7 +1159,8 @@ class LlamaRotaryEmbedding(torch.nn.Module):
             # [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