Upload HyenaDNAForCausalLM

modeling_hyena.py

@@ -19,8 +19,8 @@ def fftconv(u, k, D):
     seqlen = u.shape[-1]
     fft_size = 2 * seqlen
 
-    k_f = torch.fft.rfft(k, n=fft_size) / fft_size
-    u_f = torch.fft.rfft(u.to(dtype=k.dtype), n=fft_size)
+    k_f = torch.fft.rfft(k.to(torch.float32), n=fft_size) / fft_size
+    u_f = torch.fft.rfft(u.to(dtype=torch.float32), n=fft_size)
 
     if len(u.shape) > 3: k_f = k_f.unsqueeze(1)
     y = torch.fft.irfft(u_f * k_f, n=fft_size, norm='forward')[..., :seqlen]