class_docstr: The criterion corresponding to the SimCLR loss as defined in the paper
  https://arxiv.org/abs/2002.05709.
  
  Args:
      temperature (float): the temperature to be applied on the logits
      buffer_params:
          world_size (int): total number of trainers in training
          embedding_dim (int): output dimensions of the features projects
          effective_batch_size (int): total batch size used (includes positives)
class_docstr_tok: ["The","criterion","corresponding","to","the","SimCLR","loss","as","defined","in","the","paper","https",":","\/\/arxiv.org\/abs\/2002.05709",".","Args",":","temperature","(","float",")",":","the","temperature","to","be","applied","on","the","logits","buffer_params",":","world_size","(","int",")",":","total","number","of","trainers","in","training","embedding_dim","(","int",")",":","output","dimensions","of","the","features","projects","effective_batch_size","(","int",")",":","total","batch","size","used","(","includes","positives",")"]
docstr: Calculate the loss. Operates on embeddings tensor.
docstr_tok: ["Calculate","the","loss",".","Operates","on","embeddings","tensor","."]
code: def forward(self, embedding: torch.Tensor):
      """
      Calculate the loss. Operates on embeddings tensor.
      """
      assert embedding.ndim == 2
      assert embedding.shape[1] == int(self.buffer_params.embedding_dim)
  
      batch_size = embedding.shape[0]
      T = self.temperature
      num_pos = self.num_pos
      assert (
          batch_size % num_pos == 0
      ), "Batch size should be divisible by num_pos"
  
      # Step 1: gather all the embeddings. Shape example: 4096 x 128
      embeddings_buffer = self.gather_embeddings(embedding)
  
      # Step 2: matrix multiply: 64 x 128 with 4096 x 128 = 64 x 4096 and
      # divide by temperature.
      similarity = torch.exp(
          torch.mm(embedding, embeddings_buffer.t()) / T
      )
      pos = torch.sum(similarity * self.pos_mask, 1)
      neg = torch.sum(similarity * self.neg_mask, 1)
      loss = -(torch.mean(torch.log(pos / (pos + neg))))
      return loss
  
imports: import logging
  import pprint
  import numpy
  import torch
  from classy_vision.generic.distributed_util import get_cuda_device_index, get_rank
  from classy_vision.losses import ClassyLoss, register_loss
  from torch import nn
  from vissl.config import AttrDict
  from vissl.utils.distributed_utils import gather_from_all
label_desc: Use image node_id 3 for calling the SimclrInfoNCECriterion obj's underlying member method code with example usage: obj.forward(embedding) and returns: loss