Update modeling.py

modeling.py

@@ -4,6 +4,7 @@ import torch.nn.functional as F
 from torch.nn.attention import sdpa_kernel, SDPBackend
 from transformers import PreTrainedModel
 from .configuration import CustomConfig
+from transformers.modeling_outputs import CausalLMOutput
 
 
 class RotaryPositionalEncoding(nn.Module):
@@ -149,4 +150,5 @@ class CustomModel(PreTrainedModel):
 
     def forward(self, tensor):
         with torch.autocast('cuda', dtype=torch.bfloat16):
-            return self.model(tensor)
+            logits = self.model(tensor)
+            return CausalLMOutput(logits=logits)