How to add the original student model loss to the model distillation demo

[MeteorQ]

Hello, I have tried to use the distillation script you provided（examples/llm_distill/main.py） to distill the LLM. However, I found that only the soft label is used in the script.
I would like to ask how to modify the code so that it can be combined with the original loss (hard label) of the student model.
It would be of great help to me if you could reply! Thank you.

[coderabbitai]

📝 CodeRabbit Plan Mode

Generate an implementation plan and prompts that you can use with your favorite coding agent.

• Create Plan

Examples

• Example 1
• Example 2

---

---

🧪 Issue enrichment is currently in open beta.

You can configure auto-planning by selecting labels in the issue_enrichment configuration.

To disable automatic issue enrichment, add the following to your .coderabbit.yaml:

issue_enrichment:
  auto_enrich:
    enabled: false

💬 Have feedback or questions? Drop into our discord or schedule a call!

[AAnoosheh]

Hi there, you can see an example of passing in student_loss=loss into compute_kd_loss() here:
https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/llm_distill#distill-during-training

You also will need to use the existing basic loss balancer class we have, or write your own if you want to further customize how the losses are weighed together. See the part right above it for adding the loss balancer to the KD configuration:
https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/llm_distill#set-up-the-meta-model