code: def forward_loss(args, ds_config, output_mode):
      assert args.distill_method in ["zero_stage", "one_stage"]
  
      if args.distill_method == "zero_stage":
  
          def _simple_function(batch, model, teacher_model=None):
              outputs = model(**batch)
              return [outputs.loss, 0, 0, 0]
  
          return _simple_function
  
      elif args.distill_method == "one_stage":
          loss_mse = MSELoss()
          if output_mode == "classification":
              cls_loss_func = soft_cross_entropy
          elif output_mode == "regression":
              cls_loss_func = loss_mse
  
          def _kd_function(batch, model, teacher_model):
              (
                  att_loss,
                  rep_loss,
                  loss,
              ) = (
                  0.0,
                  0.0,
                  0.0,
              )
              outputs = model(
                  **batch,
                  output_attentions=True,
                  output_hidden_states=True,
              )
              student_logits, student_reps, student_atts = (
                  outputs.logits,
                  outputs.hidden_states,
                  outputs.attentions,
              )
              with torch.no_grad():
                  outputs_teacher = teacher_model(
                      **batch,
                      output_attentions=True,
                      output_hidden_states=True,
                  )
              teacher_logits, teacher_reps, teacher_atts = (
                  outputs_teacher.logits,
                  outputs_teacher.hidden_states,
                  outputs_teacher.attentions,
              )
              cls_loss = cls_loss_func(student_logits, teacher_logits)
              loss += cls_loss
              teacher_layer_num, student_layer_num = len(
                  teacher_atts
              ), len(student_atts)
              if args.layer_reduction_enabled:
                  teacher_layers = [
                      x
                      for x in ds_config["compression_training"][
                          "layer_reduction"
                      ]["teacher_layer"]
                  ]
                  att_list = [x for x in teacher_layers]
                  rep_list = [
                      teacher_layers[0] - 1,
                  ] + [x + 1 for x in teacher_layers]
              else:
                  ## ATTENTION: The knowledge distillation is designed for skip-layer KD
                  layers_per_block = int(
                      teacher_layer_num / student_layer_num
                  )  ###[1, 3, 5, 7, 9, 11]
                  att_list = [
                      i * layers_per_block + layers_per_block - 1
                      for i in range(student_layer_num)
                  ]
                  rep_list = [
                      i * layers_per_block
                      for i in range(student_layer_num + 1)
                  ]  ###[0, 2, 4, 6, 8, 10, 12]
  
              new_teacher_reps = [teacher_reps[i] for i in rep_list]
              new_teacher_atts = [teacher_atts[i] for i in att_list]
  
              for student_att, teacher_att in zip(
                  student_atts, new_teacher_atts
              ):
                  tmp_loss = loss_mse(student_att, teacher_att)
                  att_loss += tmp_loss
  
              for student_rep, teacher_rep in zip(
                  student_reps, new_teacher_reps
              ):
                  tmp_loss = loss_mse(student_rep, teacher_rep)
                  rep_loss += tmp_loss
              loss += att_loss + rep_loss
              return [
                  loss,
                  rep_loss.item(),
                  cls_loss.item(),
                  att_loss.item(),
              ]
  
          return _kd_function
  
imports: import torch
  from torch.nn import CrossEntropyLoss, MSELoss
  import datasets
  from datasets import load_dataset, load_metric
  import copy
  import os
  import json
  import transformers
  from transformers import AutoConfig, PretrainedConfig
  import huggingface_transformer
  from huggingface_transformer.modeling_bert import BertForSequenceClassification
  import logging
  import numpy
  import math
label_desc: Use image node_id 9 for calling a global function with example usage: forward_loss(args, ds_config, output_mode) and returns: _simple_function, list, _kd_function, list