class_docstr: This loss allows to distill a SwAV model
class_docstr_tok: ["This","loss","allows","to","distill","a","SwAV","model"]
docstr: Compute the teacher assignments toward which the student crops
  scores should converge
docstr_tok: ["Compute","the","teacher","assignments","toward","which","the","student","crops","scores","should","converge"]
code: def compute_teacher_assignements(self, teacher_scores: torch.Tensor):
      """
      Compute the teacher assignments toward which the student crops
      scores should converge
      """
  
      # In case we do not want to use Sink Horn Knopp and take the
      # teacher assignments as they are, return the softmax of
      # the teacher scores
      if self.nmb_sinkhornknopp_iters == 0:
          return torch.nn.functional.softmax(
              teacher_scores / self.epsilon, dim=-1
          )
  
      # Add teacher scores of the past stored in the queue
      batch_size = teacher_scores.shape[0]
      if self.use_queue:
          queue = self.local_queue[: self.local_queue_length]
          teacher_scores = torch.cat((teacher_scores, queue))
  
      # Divide by epsilon (which can be seen as a temperature which
      # helps to sharpen the distribution of the assignments)
      if self.use_double_prec:
          assignments = torch.exp(
              teacher_scores.double() / np.float64(self.epsilon)
          ).t()
          assignments = assignments.double()
      else:
          assignments = teacher_scores / self.epsilon
          # use the log-sum-exp trick for numerical stability.
          M = torch.max(assignments)
          all_reduce_max(M)
          assignments -= M
          assignments = torch.exp(assignments).t()
  
      # Apply sinkhornknopp algorithm to divide equally the
      # assignment to each of the prototypes
      assignments = distributed_sinkhornknopp(
          Q=assignments,
          hard_assignment=self.num_iteration
          < self.temp_hard_assignment_iters,
          world_size=self.world_size,
          num_iter=self.nmb_sinkhornknopp_iters,
          use_gpu=self.use_gpu,
          use_double_prec=self.use_double_prec,
      )
  
      # Extract the assignments of the current batch (discard queue
      # assignments used to stabilize sink-horn-knopp)
      assignments = assignments[:batch_size]
      return assignments
  
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 6 for calling the SwAVDistillationCriterion obj's underlying member method code with example usage: obj.compute_teacher_assignements(teacher_scores) and returns: assignments, torch