make sure rotary positional embeddings are same type as queries

lucidrains · Nov 6, 2021 · 3f6c461 · 3f6c461
1 parent 48a1535
commit 3f6c461
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/routing_transformer/routing_transformer.py b/routing_transformer/routing_transformer.py
@@ -236,6 +236,7 @@ def rotate_every_two(x):
     return rearrange(x, '... d j -> ... (d j)')
 
 def apply_rotary_pos_emb(q, k, v, sinu_pos):
+    sinu_pos = sinu_pos.type(q.dtype)
     sinu_pos = rearrange(sinu_pos, '() n (j d) -> n j d', j = 2)
     sin, cos = sinu_pos.unbind(dim = -2)
     sin, cos = map(lambda t: repeat(t, 'b n -> b (n j)', j = 2), (sin, cos))

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'routing_transformer',
   packages = find_packages(exclude=['examples']),
-  version = '1.6.0',
+  version = '1.6.1',
   license='MIT',
   description = 'Routing Transformer (Pytorch)',
   author = 'Phil Wang, Aran Komatsuzaki',