fixing arange not knowing half on cpu

facebookresearch · Jun 4, 2022 · 66fa7e9 · 66fa7e9
1 parent 0ef428a
commit 66fa7e9
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/tests/test_rotary_embeddings.py b/tests/test_rotary_embeddings.py
@@ -68,6 +68,9 @@ def test_rotary_embeddings(device, dtype):
     assert k_rot.dtype == k.dtype
 
     # Check that the sequences now encode relative position information
+    q, k = q.float(), k.float()
+    q_rot, k_rot = q_rot.float(), k_rot.float()
+
     att = torch.einsum("bhne,bhme->bhnm", q, k)
     att_rot = torch.einsum("bhne,bhme->bhnm", q_rot, k_rot)
 

diff --git a/xformers/components/positional_embedding/rotary.py b/xformers/components/positional_embedding/rotary.py
@@ -67,12 +67,14 @@ def _update_cos_sin_tables(self, x, seq_dimension=1):
             or self._cos_cached.dtype != x.dtype
         ):
             self._seq_len_cached = seq_len
-            t = torch.arange(x.shape[seq_dimension], device=x.device, dtype=x.dtype)
+            t = torch.arange(
+                x.shape[seq_dimension], device=x.device, dtype=torch.float32
+            )
             freqs = torch.einsum("i,j->ij", t, self.inv_freq.to(x.dtype))
             emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
 
-            self._cos_cached = emb.cos()[None, None, :, :]
-            self._sin_cached = emb.sin()[None, None, :, :]
+            self._cos_cached = emb.cos()[None, None, :, :].to(x.dtype)
+            self._sin_cached = emb.sin()[None, None, :, :].to(x.dtype)
 
         return self._cos_cached, self._sin_cached