Update modeling_auristream.py

modeling_auristream.py

@@ -529,18 +529,10 @@ class MLP(nn.Module):
 
 
 class Rotary(torch.nn.Module):
-    def __init__(self, dim, base=500000, learned=True):
+    def __init__(self, dim, base=10000, learned=False):
         super().__init__()
         # Compute the base inverse frequencies as before.
-        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
-        # If learned is True, register as a parameter; otherwise, as a buffer.
-        if learned:
-            # Initialize randomly and register as a parameter.
-            self.inv_freq = torch.nn.Parameter(inv_freq)
-            nn.init.normal_(self.inv_freq, mean=0.0, std=0.02)
-        else:
-            self.register_buffer("inv_freq", inv_freq)
-        self.learned = learned  # (optional) Save the flag if needed later
+        self.inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
 
     def forward(self, x):
         seq_len = x.shape[1]
@@ -552,6 +544,7 @@ class Rotary(torch.nn.Module):
         sin_cached = freqs.sin()
         return cos_cached[None, :, None, :], sin_cached[None, :, None, :]
 
+
 def apply_rotary_emb(x, cos, sin):
     assert x.ndim == 4  # multihead attention expected
     d = x.shape[3] // 2