fixing arange not knowing half on cpu

facebookresearch · Jun 4, 2022 · 50f6da4 · 50f6da4
1 parent 0ef428a
commit 50f6da4
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/xformers/components/positional_embedding/rotary.py b/xformers/components/positional_embedding/rotary.py
@@ -67,7 +67,9 @@ 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
+            ).to(x.dtype)
             freqs = torch.einsum("i,j->ij", t, self.inv_freq.to(x.dtype))
             emb = torch.cat((freqs, freqs), dim=-1).to(x.device)