class_docstr: Loss used in iBOT (https://arxiv.org/pdf/2111.07832.pdf)
  
  Mostly copy-pasted from the official iBOT repository:
  https://github.com/bytedance/ibot/blob/da316d82636a7a7356835ef224b13d5f3ace0489/main_ibot.py#L478
docstr: Cross-entropy between softmax outputs of the teacher and student networks.
  
  Args:
      student_output:     global views class token + feature map
      teacher_output:     global views class token + feature map
      student_local_cls:  local views class token only
      student_mask:       mask used for the global view
      epoch (int):        current epoch of training
docstr_tok: ["Cross-entropy","between","softmax","outputs","of","the","teacher","and","student","networks",".","Args",":","student_output",":","global","views","class","token","+","feature","map","teacher_output",":","global","views","class","token","+","feature","map","student_local_cls",":","local","views","class","token","only","student_mask",":","mask","used","for","the","global","view","epoch","(","int",")",":","current","epoch","of","training"]
code: def forward(
      self,
      student_output: torch.Tensor,
      teacher_output: torch.Tensor,
      student_local_cls: torch.Tensor,
      student_mask: List[torch.Tensor],
      epoch: int,
  ):
      """
      Cross-entropy between softmax outputs of the teacher and student networks.
  
      Args:
          student_output:     global views class token + feature map
          teacher_output:     global views class token + feature map
          student_local_cls:  local views class token only
          student_mask:       mask used for the global view
          epoch (int):        current epoch of training
      """
      student_cls, student_patch = student_output
      teacher_cls, teacher_patch = teacher_output
  
      if student_local_cls is not None:
          student_cls = torch.cat([student_cls, student_local_cls])
  
      # [CLS] and patch for global patches
      student_cls = student_cls / self.student_temp
      student_cls_c = student_cls.chunk(self.ncrops)
      student_patch = student_patch / self.student_temp
      student_patch_c = student_patch.chunk(self.ngcrops)
  
      # teacher centering and sharpening
      temp = self.teacher_temp_schedule[epoch]
      temp2 = self.teacher_temp2_schedule[epoch]
      teacher_cls_c = F.softmax(
          (teacher_cls - self.center) / temp, dim=-1
      )
      teacher_cls_c = teacher_cls_c.detach().chunk(self.ngcrops)
      teacher_patch_c = F.softmax(
          (teacher_patch - self.center2) / temp2, dim=-1
      )
      teacher_patch_c = teacher_patch_c.detach().chunk(self.ngcrops)
  
      total_loss1, n_loss_terms1 = 0, 0
      total_loss2, n_loss_terms2 = 0, 0
      for q in range(len(teacher_cls_c)):
          for v in range(len(student_cls_c)):
              # If same view: don't distill the class token but the patch tokens
              # for the masked image modeling criterion part of iBOT
              if v == q:
                  loss2 = torch.sum(
                      -teacher_patch_c[q]
                      * F.log_softmax(student_patch_c[v], dim=-1),
                      dim=-1,
                  )
                  mask = student_mask[v].flatten(-2, -1)
                  loss2 = torch.sum(
                      loss2 * mask.float(), dim=-1
                  ) / mask.sum(dim=-1).clamp(min=1.0)
                  total_loss2 += loss2.mean()
                  n_loss_terms2 += 1
              # If different views: distill the class tokens as
              # it encourages to push toward recognizing same content
              else:
                  loss1 = torch.sum(
                      -teacher_cls_c[q]
                      * F.log_softmax(student_cls_c[v], dim=-1),
                      dim=-1,
                  )
                  total_loss1 += loss1.mean()
                  n_loss_terms1 += 1
  
      total_loss1 = total_loss1 / n_loss_terms1 * self.lambda1
      total_loss2 = total_loss2 / n_loss_terms2 * self.lambda2
      total_loss = {
          "cls": total_loss1,
          "patch": total_loss2,
          "loss": total_loss1 + total_loss2,
      }
      self.update_center(teacher_cls, teacher_patch)
      return total_loss
  
imports: import logging
  import pprint
  from typing import List
  import numpy
  import torch
  import torch.distributed
  import torch.nn
  import torch.nn.functional
  from classy_vision.losses import ClassyLoss, register_loss
  from fairscale.nn import FullyShardedDataParallel
  from vissl.config import AttrDict
  from vissl.models.model_helpers import get_no_ddp_model
  from vissl.utils.distributed_utils import concat_all_gather
label_desc: Use image node_id 2 for calling the iBOTCriterion obj's underlying member method code with example usage: obj.forward(student_output, teacher_output, student_local_cls, student_mask, epoch) and returns: total_loss