Update modeling_auristream.py

modeling_auristream.py

@@ -72,7 +72,9 @@ class AuriStream(PreTrainedModel):
         elif isinstance(module, nn.Embedding):
             torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
    
-    def forward(self, input_ids, tgt=None, output_logits=False, output_hidden_states=False, return_dict=False, up_until_layer=None):
+    def forward(self, input_ids, tgt=None, output_logits=False, 
+                output_hidden_states=False, return_dict=False, 
+                up_until_layer=None, **kwargs):
         """
         Input: coch: torch.Tensor of shape (b, t)
                tgt_coch: torch.Tensor of shape (b, t) or None