Adding modeling.py file

modeling.py

@@ -109,7 +109,7 @@ class MoLM(PreTrainedModel):
                 # Expand the expert_mask to match the logits shape (batch_size, 1, 1)
                 expert_mask_expanded = expert_mask.unsqueeze(-1).unsqueeze(-1).float()
 
-                expert_output = expert(input_ids, targets=targets, date=date, **kwargs)
+                expert_output = expert(input_ids, targets=targets, date=date, **kwargs, get_logits=True)
 
                 logits = expert_output["logits"]
                 loss_to_log = expert_output["loss_to_log"]