docstr: The second return argument is used for unit testing.
  
  Args:
      * sample: batched sample that has teacher score keys (top_k_scores and
       top_k_indices)
      * student_lprobs: tensor of student log probabilities
      * lprobs: flat version of student_lprobs
docstr_tok: ["The","second","return","argument","is","used","for","unit","testing",".","Args",":","*","sample",":","batched","sample","that","has","teacher","score","keys","(","top_k_scores","and","top_k_indices",")","*","student_lprobs",":","tensor","of","student","log","probabilities","*","lprobs",":","flat","version","of","student_lprobs"]
code: def get_kd_loss(self, sample, student_lprobs, lprobs):
      """
      The second return argument is used for unit testing.
  
      Args:
          * sample: batched sample that has teacher score keys (top_k_scores and
           top_k_indices)
          * student_lprobs: tensor of student log probabilities
          * lprobs: flat version of student_lprobs
      """
      top_k_teacher_probs_normalized = sample["top_k_scores"]
      indices = sample["top_k_indices"]
  
      assert indices.shape[0:1] == student_lprobs.shape[0:1]
  
      kd_loss = -(
          torch.sum(
              torch.gather(student_lprobs, 2, indices)
              * top_k_teacher_probs_normalized.float()
          )
      )
      return kd_loss
  
file_path: translate/pytorch_translate/research/knowledge_distillation/knowledge_distillation_loss.py
imports: import math
  import torch
  import torch.nn.functional
  from fairseq import utils
  from fairseq.criterions import LegacyFairseqCriterion, register_criterion
  from pytorch_translate import utils
label_desc: Use image node_id 3 for calling the KnowledgeDistillationCriterion obj's underlying member method code with example usage: obj.get_kd_loss(sample, student_lprobs, lprobs) and returns: kd_loss