Update modeling_my_gpt.py

modeling_my_gpt.py

@@ -1,6 +1,8 @@
 import torch
 import torch.nn as nn
 from transformers import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutput
+
 from .configuration_my_gpt import MyGPTConfig
 from .untrained_model import GPTModel
 
@@ -32,9 +34,9 @@ class MyGPTForCausalLM(PreTrainedModel):
 
         self.post_init()
 
-    def forward(self, input_ids, labels=None):
+    def forward(self, input_ids, labels=None, **kwargs):
         logits = self.model(input_ids)
-
+    
         loss = None
         if labels is not None:
             shift_logits = logits[..., :-1, :].contiguous()
@@ -44,8 +46,9 @@ class MyGPTForCausalLM(PreTrainedModel):
                 shift_logits.view(-1, shift_logits.size(-1)),
                 shift_labels.view(-1)
             )
+    
+        return CausalLMOutput(
+            loss=loss,
+            logits=logits,
+        )
 
-        return {
-            "loss": loss,
-            "logits": logits,
-        }