Update modeling_auristream.py

modeling_auristream.py

@@ -532,7 +532,9 @@ class Rotary(torch.nn.Module):
     def __init__(self, dim, base=10000, learned=False):
         super().__init__()
         # Compute the base inverse frequencies as before.
-        self.inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        # self.inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
 
     def forward(self, x):
         seq_len = x.shape[1]