ix

layers.py

@@ -57,7 +57,6 @@ class STULayer(nn.Module):
     
         # Normalize and apply STU
         x_normed = self.stu_norm(x).to(dtype=self.stu.M_inputs.dtype)  # Match dtype for STU
-        print(f"x dtype: {x.dtype}")
         x_stu = self.stu(x_normed).to(dtype=x.dtype)  # Ensure output matches `x`'s dtype
         x = x + x_stu
     

stu.py

@@ -29,11 +29,12 @@ class STU(nn.Module):
         self.d_out = config.n_embd
         self.use_hankel_L = config.use_hankel_L
         self.use_approx = config.use_approx
-        self.flash_fft = (
-            FlashFFTConv(self.n, dtype=torch_dtype)
-            if config.use_flash_fft and flash_fft_available
-            else None
-        )
+        self.flash_fft = None
+        if config.use_flash_fft and flash_fft_available:
+            if torch_dtype == torch.float16:  # Only enable for float16
+                self.flash_fft = FlashFFTConv(self.n, dtype=torch.float16)
+        else:
+            print(f"Disabling FlashFFTConv for unsupported dtype: {torch_dtype}")
         if self.use_approx:
             self.M_inputs = nn.Parameter(
                 torch.empty(self.d_in, self.d_out, dtype=torch_dtype)