class_docstr: This loss allows to distill a SwAV model
class_docstr_tok: ["This","loss","allows","to","distill","a","SwAV","model"]
code: def forward(
      self,
      logits: torch.Tensor,
      teacher_scores: torch.Tensor,
      teacher_prototypes: Optional[torch.Tensor],
  ):
      # Find the batch size and make sure inputs are correct
      assert logits.shape[0] % self.num_crops == 0
      batch_size = logits.shape[0] // self.num_crops
      assert teacher_scores.shape[0] % batch_size == 0
      num_crops_for_assign = teacher_scores.shape[0] // batch_size
  
      # Compute the scores based on the teacher prototypes
      if teacher_prototypes is not None:
          assert not teacher_prototypes.requires_grad
          scores = functional.linear(logits, teacher_prototypes)
  
      # Or take the logits of the student directly as scores if
      # the teacher prototypes are not provided as input
      else:
          message = f"Incompatible feature dimensions {logits.shape[1]} vs {teacher_scores.shape[1]}"
          assert logits.shape[1] == teacher_scores.shape[1], message
          scores = logits
  
      # Compute the loss toward each of the assignments
      loss = 0
      for i in range(num_crops_for_assign):
          teacher_crop_scores = teacher_scores[
              i * batch_size : (i + 1) * batch_size
          ]
  
          # If just one crop to assign, decide whether or not we should map
          # the two crops of the student or the other one to the teacher
          if num_crops_for_assign == 1:
              start_crop = 1 if self.swapped_assignment else 0
              student_crop_ids = list(range(start_crop, self.num_crops))
  
          # If several crops to assign, either use swapped assignments or
          # non-swapped assignments
          elif self.swapped_assignment:
              student_crop_ids = [
                  p for p in range(self.num_crops) if p != i
              ]
          else:
              student_crop_ids = [i] + list(range(2, self.num_crops))
  
          loss += self.forward_to_teacher_crop(
              scores,
              teacher_crop_scores,
              batch_size=batch_size,
              student_crop_ids=student_crop_ids,
          )
      return loss / num_crops_for_assign
  
imports: import logging
  import os
  import pprint
  from typing import List, Optional, Union
  import numpy
  import torch
  import torch.nn.functional
  from classy_vision.generic.distributed_util import all_reduce_max, get_cuda_device_index, get_rank, get_world_size
  from classy_vision.losses import ClassyLoss, register_loss
  from fvcore.common.file_io import PathManager
  from torch import nn
  from vissl.config import AttrDict
  from vissl.losses.distibuted_sinkhornknopp import distributed_sinkhornknopp
label_desc: Use image node_id 4 for calling the SwAVDistillationCriterion obj's underlying member method code with example usage: obj.forward(logits, teacher_scores, teacher_prototypes) and returns: unknown