Update modeling.py

modeling.py

@@ -150,7 +150,7 @@ class RecombinationTransformerForCausalLM(PreTrainedModel):
         self.final_rms_norm = RMSNorm(config.embed_dim)
         self.lm_head = nn.Linear(config.embed_dim, config.vocab_size, bias=False)
     
-    def forward(self, input_ids, attention_mask=None, past_key_values=None):
+    def forward(self, input_ids, attention_mask=None, past_key_values=None, return_dict=None):
         if attention_mask is None:
             attention_mask = torch.ones(input_ids.shape, device=input_ids.device)
         
@@ -171,6 +171,9 @@ class RecombinationTransformerForCausalLM(PreTrainedModel):
         # LM head
         logits = self.lm_head(x)
         
+        if not return_dict:
+            return (logits,)
+
         return CausalLMOutputWithPast(logits=logits, past_key_values=past_key_values)
         
     def prepare_inputs_for_generation(self, input_ids, past=None, attention_mask=None, **kwargs):