only ipu cos sin in fp16

graphcore-research · Oct 20, 2023 · 40437d0 · 40437d0
1 parent dc77144
commit 40437d0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/besskge/utils.py b/besskge/utils.py
@@ -66,7 +66,7 @@ def complex_rotation(v: torch.Tensor, r: torch.Tensor) -> torch.Tensor:
         Row-wise rotated tensors.
     """
     # Always compute sin and cos in fp16, as faster on IPU
-    if r.dtype == torch.float32:
+    if r.dtype == torch.float32 and r.device.type == "ipu":
         r_cos = torch.cos(r.to(dtype=torch.float16)).to(dtype=torch.float32)
         r_sin = torch.sin(r.to(dtype=torch.float16)).to(dtype=torch.float32)
     else: