Update modeling_minGRULM.py

modeling_minGRULM.py

@@ -99,16 +99,15 @@ class MinGRULMForCausalLM(PreTrainedModel):
         # Ensure that inputs for generation are properly handled
         return {"input_ids": input_ids, "attention_mask": kwargs.get("attention_mask", None)}
 
-    def forward(
-        self,
-        input_ids: torch.LongTensor,
-        labels: Optional[torch.LongTensor] = None,
-        return_dict: Optional[bool] = True,
-        **kwargs
-    ):
+    def forward(self, input_ids: torch.LongTensor, labels: Optional[torch.LongTensor] = None, return_dict: Optional[bool] = True, **kwargs):
         # Forward pass through the wrapped model
         logits = self.model(input_ids)
 
+        # NaN kontrolü: Eğer logits'te NaN varsa, sıfırlama
+        if torch.isnan(logits).any():
+            print("NaN detected in logits! Replacing with zeros.")
+            logits = torch.nan_to_num(logits, nan=0.0)
+
         loss = None
         if labels is not None:
             shift_logits = logits[..., :-1, :].contiguous()
@@ -119,6 +118,11 @@ class MinGRULMForCausalLM(PreTrainedModel):
                 shift_labels.view(-1),
             )
 
+            # NaN kontrolü: Eğer loss'ta NaN varsa, sıfırlama
+            if torch.isnan(loss).any():
+                print("NaN detected in loss! Replacing with zeros.")
+                loss = torch.tensor(0.0, device=loss.device)  # NaN olan loss'u sıfırlıyoruz
+
         if not return_dict:
             return (loss, logits) if loss is not None else (logits,)