Patch Fix

attn.py

@@ -52,14 +52,15 @@ class Attention(nn.Module):
             # Get slopes for the nearest power of two
             n = nearest_power_of_two(n_heads, round_up=False)
             slopes_power_of_two = self._generate_slopes(n)
-
+    
             # Generate extra slopes
             extra_slopes = self._generate_slopes(2 * n)
             extra_slopes_trunc = extra_slopes[0::2][: n_heads - n]
             slopes = slopes_power_of_two + extra_slopes_trunc
         slopes = torch.tensor(slopes, device=self.device)
         slopes = slopes * interpolation_factor  # https://arxiv.org/pdf/2310.13017
-        return slopes
+        return slopes.to(torch.float32)  # Ensure slopes are in float32
+
 
     def forward(self, x):
         bsz, seq_len, d_in = x.size()