code: def _get_kl_divergence_penalty(
      self, student_logits: torch.Tensor, teacher_logits: torch.Tensor
  ):
      # Computes KL(teacher || student) not KL(student || teacher)
      # Note: PyTorch KL Divergence expects log probabilities
      student_log_probs = functional.log_softmax(
          student_logits / self.temperature, dim=-1
      )
      teacher_probs = functional.softmax(
          teacher_logits / self.temperature, dim=-1
      )
      return functional.kl_div(
          student_log_probs,
          teacher_probs,
          reduction="batchmean",
          log_target=False,
      )
  
imports: from enum import Enum
  import torch
  import torch.nn
  import torch.nn.functional
  from classy_vision.losses import ClassyLoss, register_loss
  from vissl.config.attr_dict import AttrDict
  from vissl.losses.cross_entropy_multiple_output_single_target import SmoothCrossEntropy
label_desc: Use image node_id 4 for calling the DistillationCriterion obj's underlying member method code with example usage: obj._get_kl_divergence_penalty(student_logits, teacher_logits) and returns: functional