memories should not have relative positions with each other, just far…

… away from main tokens
lucidrains · Oct 2, 2023 · 974f43d · 974f43d
1 parent 9261e86
commit 974f43d
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'voicebox-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.2.2',
+  version = '0.2.3',
   license='MIT',
   description = 'Voicebox - Pytorch',
   author = 'Phil Wang',

diff --git a/voicebox_pytorch/voicebox_pytorch.py b/voicebox_pytorch/voicebox_pytorch.py
@@ -383,10 +383,12 @@ def forward(
 
         # rotary embeddings
 
-        main_positions = torch.arange(seq_len, device = self.device, dtype = torch.long)
-        register_positions = torch.arange(self.num_register_tokens, device = self.device, dtype = torch.long)
-        register_positions -= 10000
-        positions = torch.cat((register_positions, main_positions))
+        positions = seq_len
+
+        if self.has_register_tokens:
+            main_positions = torch.arange(seq_len, device = self.device, dtype = torch.long)
+            register_positions = torch.full((self.num_register_tokens,), -10000, device = self.device, dtype = torch.long)
+            positions = torch.cat((register_positions, main_positions))
 
         rotary_emb = self.rotary_emb(positions)