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: Do a gather over all embeddings, so we can compute the loss.
  Final shape is like: (batch_size * num_gpus) x embedding_dim
docstr_tok: ["Do","a","gather","over","all","embeddings",",","so","we","can","compute","the","loss",".","Final","shape","is","like",":","(","batch_size","*","num_gpus",")","x","embedding_dim"]
code: def gather_embeddings(embedding: torch.Tensor):
      """
      Do a gather over all embeddings, so we can compute the loss.
      Final shape is like: (batch_size * num_gpus) x embedding_dim
      """
      if (
          torch.distributed.is_available()
          and torch.distributed.is_initialized()
      ):
          embedding_gathered = gather_from_all(embedding)
      else:
          embedding_gathered = embedding
      return embedding_gathered
  
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 5 for calling the SimclrInfoNCECriterion obj's underlying member method code with example usage: obj.gather_embeddings(embedding) and returns: embedding_gathered