From c058d1aa96b46dd9ff4248967d4190e012ac8e09 Mon Sep 17 00:00:00 2001
From: Daniel Han-Chen <danielhanchen@gmail.com>
Date: Sun, 25 Feb 2024 02:17:45 +1100
Subject: [PATCH] Update gemma.py

---
 unsloth/models/gemma.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py
index 22734a0933..51432f337b 100644
--- a/unsloth/models/gemma.py
+++ b/unsloth/models/gemma.py
@@ -57,7 +57,7 @@ def fast_geglu_inference(self, X):
 pass
 
 
-class FastGemmaRotaryEmbedding(nn.Module):
+class FastGemmaRotaryEmbedding(torch.nn.Module):
     def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
         super().__init__()