fixes ci checks

Author: sarathc-cerebras

src/transformers/models/jais2/configuration_jais2.py

@@ -1,4 +1,6 @@
 from ...configuration_utils import PretrainedConfig
+
+
 # from ...modeling_rope_utils import rope_config_validation, standardize_rope_params
 
 
@@ -171,7 +173,6 @@ def __init__(
         self.attention_dropout = attention_dropout
         self.mlp_bias = mlp_bias
         self.head_dim = head_dim if head_dim is not None else self.hidden_size // self.num_attention_heads
-        
         # Set up rope_parameters from rope_scaling
         self.rope_parameters = rope_scaling
 

src/transformers/models/jais2/modeling_jais2.py

@@ -918,6 +918,7 @@ def forward(
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
         cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
     ) -> Union[tuple, BaseModelOutputWithPast]:
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (
@@ -1151,7 +1152,6 @@ def _prepare_4d_causal_attention_mask_with_cache_position(
 
 class Jais2ForCausalLM(Jais2PreTrainedModel, GenerationMixin):
     _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
-    # _tied_weights_keys = ["lm_head.weight"]
 
     def __init__(self, config):
         super().__init__(config)
@@ -1321,6 +1321,7 @@ def forward(
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
+        **kwargs,
     ) -> Union[tuple, SequenceClassifierOutputWithPast]:
         r"""
         labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
@@ -1366,7 +1367,7 @@ def forward(
 
         loss = None
         if labels is not None:
-            loss = self.loss_function(logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config)
+            loss = self.loss_function(logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config, **kwargs)
 
         if not return_dict:
             output = (pooled_logits,) + transformer_outputs[1:]
@@ -1516,6 +1517,7 @@ def forward(
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
+        **kwargs,
     ) -> Union[tuple, TokenClassifierOutput]:
         r"""
         labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
@@ -1542,7 +1544,7 @@ def forward(
 
         loss = None
         if labels is not None:
-            loss = self.loss_function(logits=logits, labels=labels, config=self.config)
+            loss = self.loss_function(logits=logits, labels=labels, config=self.config, **kwargs)
 
         if not return_dict:
             output = (logits,) + outputs[2:]