fix

attn.py

@@ -37,7 +37,6 @@ def _get_alibi_slopes(self, n_heads: int, interpolation_factor: float = 0.25):
         slopes = slopes_power_of_two + extra_slopes_trunc
     slopes = torch.tensor(slopes, device=self.device, dtype=torch.float32)
     slopes = slopes * interpolation_factor  # https://arxiv.org/pdf/2310.13017
-    slopes = slopes.to(torch.float32)
     return slopes
 
 
@@ -127,7 +126,7 @@ class Attention(nn.Module):
             extra_slopes = self._generate_slopes(2 * n)
             extra_slopes_trunc = extra_slopes[0::2][: num_heads - n]
             slopes = slopes_power_of_two + extra_slopes_trunc
-        slopes = torch.tensor(slopes, device=torch.device("cuda"))
+        slopes = torch.tensor(slopes, device=torch.device("cuda"), dtype=torch.float32)
         slopes = slopes * interpolation_factor  # https://arxiv.org/pdf/2310.13017
         return slopes