Update modeling_auristream.py

modeling_auristream.py

@@ -72,18 +72,18 @@ 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=seq, 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):
         """
         Input: coch: torch.Tensor of shape (b, t)
                tgt_coch: torch.Tensor of shape (b, t) or None
         """
 
         # forward the GPT model itself
-        tok_emb = self.transformer.wte(seq) # token embeddings of shape (b, t, n_embd)
+        tok_emb = self.transformer.wte(input_ids) # token embeddings of shape (b, t, n_embd)
 
         # if wpe exists in self.transformer apply leanred positional embedding
         if hasattr(self.transformer, 'wpe'):
-            pos = torch.arange(0, seq.size(1), dtype=torch.long, device=seq.device)
+            pos = torch.arange(0, input_ids.size(1), dtype=torch.long, device=input_ids.device)
             pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd)
             x = self.transformer.drop(tok_emb + pos_emb)
         else: