Update modeling_llama3.py

modeling_llama3.py

@@ -341,7 +341,7 @@ class Llama3ForCausalLM(Llama3PreTrainedModel, GenerationMixin):
         self.vocab_size = self.text_config.vocab_size
         self.model = Llama3TextModel._from_config(config, attn_implementation=config._attn_implementation)
         self.lm_head = nn.Linear(self.text_config.hidden_size, self.vocab_size, bias=False)
-
+        self.loss_function = ForCausalLMLoss
         self.post_init()
 
     def get_input_embeddings(self):
@@ -409,7 +409,7 @@ class Llama3ForCausalLM(Llama3PreTrainedModel, GenerationMixin):
 
         loss = None
         if labels is not None:
-            loss = ForCausalLMLoss(logits, labels, self.vocab_size, **loss_kwargs)
+            loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs)
 
         if not return_dict:
             output = (logits,) + outputs[1:]