class_docstr: Teacher-Student knowledge distillation helper.
  Use this object when training a model with KD and a teacher model.
  
  Args:
      teacher_model (TrainableModel): teacher model
      temperature (float, optional): KD temperature. Defaults to 1.0.
      dist_w (float, optional): distillation loss weight. Defaults to 0.1.
      loss_w (float, optional): student loss weight. Defaults to 1.0.
      loss_function (str, optional): loss function to use (kl for KLDivLoss,
          mse for MSELoss)
docstr: Add KD loss
  
  Args:
      loss: student loss
      student_logits: student model logits
      teacher_logits: teacher model logits
  
  Returns:
      KD loss
code: def distill_loss_dict(
      self, loss, student_logits_dict, teacher_logits_dict
  ):
      """
      Add KD loss
  
      Args:
          loss: student loss
          student_logits: student model logits
          teacher_logits: teacher model logits
  
      Returns:
          KD loss
      """
  
      student_sm_dict = {}
      for i in range(len(student_logits_dict.keys())):
          student_sm_dict[i] = F.log_softmax(
              student_logits_dict[i] / self.t, dim=-1
          )
  
      teacher_sm_dict = {}
      for i in range(len(teacher_logits_dict.keys())):
          teacher_sm_dict[i] = F.softmax(
              teacher_logits_dict[i] / self.t, dim=-1
          )
  
      distill_losses = [
          self.loss_fn(
              input=student_sm_dict[i], target=teacher_sm_dict[i]
          )
          for i in range(len(student_sm_dict.keys()))
      ]
      distill_loss = torch.mean(torch.stack(distill_losses))
  
      return (self.loss_w * loss) + (
          distill_loss * self.dist_w * (self.t**2)
      )
  
file_path: nlp-architect/nlp_architect/nn/torch/distillation.py
imports: import argparse
  import logging
  import torch
  import torch.nn
  import torch.nn.functional
  from nlp_architect.models import TrainableModel
label_desc: Use image node_id 5 for calling the TeacherStudentDistill obj's underlying member method code with example usage: obj.distill_loss_dict(loss, student_logits_dict, teacher_logits_dict) and returns: unknown