code: def criterion(
      self,
      student_scores: torch.Tensor,
      teacher_probs: torch.Tensor,
  ):
      student_out = student_scores / self.student_temperature
      student_out = student_out.chunk(self.student_num_crops)
      teacher_out = teacher_probs.chunk(self.teacher_num_crops)
  
      total_loss = 0
      n_loss_terms = 0
      for iq, q in enumerate(teacher_out):
          for v in range(len(student_out)):
              if self.swapped_assignment:
                  # We skip cases where student and teacher operate on the same view
                  # in case we do swapped assignments
                  if v == iq:
                      continue
              else:
                  # Else we skip the cross assignments between big crops
                  if v < self.teacher_num_crops and v != iq:
                      continue
  
              loss = torch.sum(
                  -q * functional.log_softmax(student_out[v], dim=-1),
                  dim=-1,
              )
              total_loss += loss.mean()
              n_loss_terms += 1
      total_loss /= n_loss_terms
      return total_loss
  
file_path: vissl/vissl/losses/msn_distillation_loss.py
imports: import pprint
  from typing import List, Optional
  import torch
  import torch.nn.functional
  from classy_vision.losses import ClassyLoss, register_loss
  from vissl.config import AttrDict
label_desc: Use image node_id 4 for calling the MSNDistillationLoss obj's underlying member method code with example usage: obj.criterion(student_scores, teacher_probs) and returns: total_loss