Update modeling_auristream.py

modeling_auristream.py

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