Correctly mask z as a buffer

modeling_hyena.py

@@ -62,8 +62,8 @@ class HyenaPositionalEmbedding(nn.Module):
         f = torch.linspace(1e-4, bands - 1, bands)[None, None]
 
         z = torch.cat([t, torch.cos(-f * w), torch.sin(-f * w)], dim=-1)
-        # The original code sets z's LR to lr_pos_emb, which is 1e-5 by default
-        self.z = nn.Parameter(z, requires_grad=True)
+
+        self.register_buffer("z", z)
         self.register_buffer("t", t)
 
     def forward(self, L):