fix rotary embedding caching

lucidrains · Apr 5, 2022 · 7889d97 · 7889d97
1 parent 05a7c46
commit 7889d97
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 4 deletions.
diff --git a/palm_pytorch/palm_pytorch.py b/palm_pytorch/palm_pytorch.py
@@ -111,7 +111,7 @@ def get_rotary_embedding(self, n, device):
             return self.pos_emb[:n]
 
         pos_emb = self.rotary_emb(n, device=device)
-        self.register_buffer("position", pos_emb, persistent=False)
+        self.register_buffer("pos_emb", pos_emb, persistent=False)
         return pos_emb
 
     def forward(self, x):

diff --git a/palm_pytorch/triton/palm.py b/palm_pytorch/triton/palm.py
@@ -95,7 +95,7 @@ def get_rotary_embedding(self, n, device):
             return self.pos_emb[:n]
 
         pos_emb = self.rotary_emb(n, device=device)
-        self.register_buffer("position", pos_emb, persistent=False)
+        self.register_buffer("pos_emb", pos_emb, persistent=False)
         return pos_emb
 
     def forward(self, x):

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
     name="PaLM-pytorch",
     packages=find_packages(exclude=[]),
-    version="0.0.11",
+    version="0.0.12",
     license="MIT",
     description="PaLM: Scaling Language Modeling with Pathways - Pytorch",
     author="Phil Wang",

diff --git a/train.py b/train.py
@@ -8,7 +8,7 @@
 from torch.nn import functional as F
 from torch.utils.data import DataLoader, Dataset
 
-from palm_pytorch.triton import PaLM
+from palm_pytorch import PaLM
 from palm_pytorch.autoregressive_wrapper import AutoregressiveWrapper
 
 # constants