fix

convolve.py

@@ -41,8 +41,6 @@ def flash_convolve(
     u: torch.Tensor, v: torch.Tensor, flash_fft: FlashFFTConv, use_approx: bool = True,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     dtype = u.dtype  # Store the original dtype
-    u = u.to(torch.float32)
-    v = v.to(torch.float32)
 
     bsz, seq_len, d_in = u.shape
     _, K = v.shape
@@ -50,7 +48,7 @@ def flash_convolve(
     padded_len = nearest_power_of_two(seq_len, round_up=True)
     pad_len = padded_len - seq_len
 
-    sgn = torch.full((1, 1, padded_len), 1, device=u.device, dtype=torch.float32)
+    sgn = torch.full((1, 1, padded_len), 1, device=u.device, dtype=dtype)
     sgn[:, :, 1::2] = -1
 
     if use_approx:

stu.py

@@ -30,7 +30,7 @@ class STU(nn.Module):
         self.use_hankel_L = config.use_hankel_L
         self.use_approx = config.use_approx
         self.flash_fft = (
-            FlashFFTConv(self.n, dtype=torch.bfloat16)
+            FlashFFTConv(self.n, dtype=torch_dtype)
             if config.use_flash_fft and flash_fft_available
             else None
         )