diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py
index 88dc0cc7bd..f8853f61d3 100644
--- a/unsloth/models/gemma.py
+++ b/unsloth/models/gemma.py
@@ -271,7 +271,7 @@ class FastGemmaModel(FastLlamaModel):
         # https://github.com/huggingface/transformers/pull/27931
         # https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/llama/modeling_llama.py
         import transformers.models.gemma.modeling_gemma
-        # transformers.models.gemma.modeling_gemma.GemmaRotaryEmbedding = LlamaRotaryEmbedding
+        transformers.models.gemma.modeling_gemma.GemmaRotaryEmbedding = LlamaRotaryEmbedding
         return
     pass
 
diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 4e0408821a..d58eb1111a 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -821,7 +821,7 @@ class LlamaRotaryEmbedding(torch.nn.Module):
         self.register_buffer("sin_cached", emb.sin().to(dtype=dtype, device=device, non_blocking=True), persistent=False)
     pass
 
-    def forward(self, x, seq_len=None):
+    def forward(self, x, position_ids, seq_len=None):
         # x: [bs, num_attention_heads, seq_len, head_size]
         if seq_len > self.max_seq_len_cached:
             self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)