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: We precompute the positive and negative masks to speed up the loss calculation
docstr_tok: ["We","precompute","the","positive","and","negative","masks","to","speed","up","the","loss","calculation"]
code: def precompute_pos_neg_mask(self):
      """
      We precompute the positive and negative masks to speed up the loss calculation
      """
      # computed once at the begining of training
      total_images = self.buffer_params.effective_batch_size
      world_size = self.buffer_params.world_size
      batch_size = total_images // world_size
      orig_images = batch_size // self.num_pos
      rank = self.dist_rank
  
      pos_mask = torch.zeros(batch_size, total_images)
      neg_mask = torch.zeros(batch_size, total_images)
  
      all_indices = np.arange(total_images)
      pos_members = orig_images * np.arange(self.num_pos)
      orig_members = torch.arange(orig_images)
      for anchor in np.arange(self.num_pos):
          for img_idx in range(orig_images):
              delete_inds = batch_size * rank + img_idx + pos_members
              neg_inds = torch.tensor(
                  np.delete(all_indices, delete_inds)
              ).long()
              neg_mask[anchor * orig_images + img_idx, neg_inds] = 1
          for pos in np.delete(np.arange(self.num_pos), anchor):
              pos_inds = (
                  batch_size * rank + pos * orig_images + orig_members
              )
              pos_mask[
                  torch.arange(
                      anchor * orig_images, (anchor + 1) * orig_images
                  ).long(),
                  pos_inds.long(),
              ] = 1
      self.pos_mask = (
          pos_mask.cuda(non_blocking=True) if self.use_gpu else pos_mask
      )
      self.neg_mask = (
          neg_mask.cuda(non_blocking=True) if self.use_gpu else neg_mask
      )
  
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 2 for calling the SimclrInfoNCECriterion obj's underlying member method code with example usage: obj.precompute_pos_neg_mask() without return types