Update modeling_auristream.py

modeling_auristream.py

@@ -293,6 +293,7 @@ class AuriStreamModel(AuriStreamPreTrainedModel):
         output_hidden_states: Optional[bool] = False,
         return_dict: Optional[bool] = True,
         # Legacy arguments for compatibility
+        return_logits: Optional[bool] = True,
         seq: Optional[torch.LongTensor] = None,
         tgt: Optional[torch.LongTensor] = None,
     ):
@@ -356,11 +357,11 @@ class AuriStreamModel(AuriStreamPreTrainedModel):
         if not return_dict:
             if labels is not None:
                 return logits, loss
-            return logits, None
+            return logits.unsuqeeze(0), None
         
         return CausalLMOutput(
             loss=loss,
-            logits=logits,
+            logits=logits.unsuqeeze(0),
             hidden_states=all_hidden_states if output_hidden_states else None,
         )