from transformers.loss.loss_utils import ForCausalLMLoss
model.loss_function = ForCausalLMLoss