class_docstr: Distributed version of the NCE loss. It performs an "all_gather" to gather
  the allocated buffers like memory no a single gpu. For this, Pytorch distributed
  backend is used. If using NCCL, one must ensure that all the buffer are on GPU.
  This class supports training using both NCE and CrossEntropy (InfoNCE).
  
  This loss is used by NPID (https://arxiv.org/pdf/1805.01978.pdf), NPID++ and
  PIRL (https://arxiv.org/abs/1912.01991) approaches.
  
  Written by: Ishan Misra (imisra@fb.com)
  
  Config params:
      norm_embedding (bool):           whether to normalize embeddings
      temperature (float):             the temperature to apply to logits
      norm_constant (int):             Z parameter in the NCEAverage
      update_mem_with_emb_index (int): In case we have multiple embeddings used
                                       in the nce loss, specify which embedding
                                       to use to update the memory.
      loss_type (str):                 options are "nce" | "cross_entropy". Using the
                                       cross_entropy turns the loss into InfoNCE loss.
      loss_weights (List[float]):      if the NCE loss is computed between multiple pairs,
                                       we can set a loss weight per term can be used to weight
                                       different pair contributions differently
      negative_sampling_params:
          num_negatives (int):         how many negatives to contrast with
          type (str):                  how to select the negatives. options "random"
      memory_params:
          memory_size (int):           number of training samples as all the samples are
                                       stored in memory
          embedding_dim (int):         the projection head output dimension
          momentum (int):              momentum to use to update the memory
          norm_init (bool):            whether to L2 normalize the initialized memory bank
          update_mem_on_forward (bool): whether to update memory on the forward pass
      num_train_samples (int):         number of unique samples in the training dataset
class_docstr_tok: ["Distributed","version","of","the","NCE","loss",".","It","performs","an","``","all_gather","''","to","gather","the","allocated","buffers","like","memory","no","a","single","gpu",".","For","this",",","Pytorch","distributed","backend","is","used",".","If","using","NCCL",",","one","must","ensure","that","all","the","buffer","are","on","GPU",".","This","class","supports","training","using","both","NCE","and","CrossEntropy","(","InfoNCE",")",".","This","loss","is","used","by","NPID","(","https",":","\/\/arxiv.org\/pdf\/1805.01978.pdf",")",",","NPID++","and","PIRL","(","https",":","\/\/arxiv.org\/abs\/1912.01991",")","approaches",".","Written","by",":","Ishan","Misra","(","imisra","@","fb.com",")","Config","params",":","norm_embedding","(","bool",")",":","whether","to","normalize","embeddings","temperature","(","float",")",":","the","temperature","to","apply","to","logits","norm_constant","(","int",")",":","Z","parameter","in","the","NCEAverage","update_mem_with_emb_index","(","int",")",":","In","case","we","have","multiple","embeddings","used","in","the","nce","loss",",","specify","which","embedding","to","use","to","update","the","memory",".","loss_type","(","str",")",":","options","are","``","nce","''","|","``","cross_entropy","''",".","Using","the","cross_entropy","turns","the","loss","into","InfoNCE","loss",".","loss_weights","(","List","[","float","]",")",":","if","the","NCE","loss","is","computed","between","multiple","pairs",",","we","can","set","a","loss","weight","per","term","can","be","used","to","weight","different","pair","contributions","differently","negative_sampling_params",":","num_negatives","(","int",")",":","how","many","negatives","to","contrast","with","type","(","str",")",":","how","to","select","the","negatives",".","options","``","random","''","memory_params",":","memory_size","(","int",")",":","number","of","training","samples","as","all","the","samples","are","stored","in","memory","embedding_dim","(","int",")",":","the","projection","head","output","dimension","momentum","(","int",")",":","momentum","to","use","to","update","the","memory","norm_init","(","bool",")",":","whether","to","L2","normalize","the","initialized","memory","bank","update_mem_on_forward","(","bool",")",":","whether","to","update","memory","on","the","forward","pass","num_train_samples","(","int",")",":","number","of","unique","samples","in","the","training","dataset"]
docstr: Sync memory across all processes before first forward pass. Only needed
  in the distributed case.
  After the first forward pass, the update_memory function in NCEAverage
  does a gather over all embeddings, so memory stays in sync. Doing a gather
  over embeddings is O(batch size). Syncing memory is O(num items in memory).
  Generally, batch size << num items in memory. So, we prefer doing the syncs
  in update_memory.
docstr_tok: ["Sync","memory","across","all","processes","before","first","forward","pass",".","Only","needed","in","the","distributed","case",".","After","the","first","forward","pass",",","the","update_memory","function","in","NCEAverage","does","a","gather","over","all","embeddings",",","so","memory","stays","in","sync",".","Doing","a","gather","over","embeddings","is","O","(","batch","size",")",".","Syncing","memory","is","O","(","num","items","in","memory",")",".","Generally",",","batch","size","<","<","num","items","in","memory",".","So",",","we","prefer","doing","the","syncs","in","update_memory","."]
code: def sync_memory(self):
      """
      Sync memory across all processes before first forward pass. Only needed
      in the distributed case.
      After the first forward pass, the update_memory function in NCEAverage
      does a gather over all embeddings, so memory stays in sync. Doing a gather
      over embeddings is O(batch size). Syncing memory is O(num items in memory).
      Generally, batch size << num items in memory. So, we prefer doing the syncs
      in update_memory.
      """
      self.nce_average.memory = all_reduce_mean(self.nce_average.memory)
      logging.info(f"Rank: {get_rank()}: Memory synced")
      # set to true once we are done. forward pass in nce_average will sync after.
      self.init_sync_memory = True
  
imports: import logging
  import math
  import pprint
  from typing import List, Union
  import numpy
  import torch
  from classy_vision.generic.distributed_util import all_reduce_mean, gather_from_all, get_rank
  from classy_vision.generic.util import is_pos_int
  from classy_vision.losses import ClassyLoss, register_loss
  from torch import nn
  from vissl.config import AttrDict
label_desc: Use image node_id 4 for calling the NCELossWithMemory obj's underlying member method code with example usage: obj.sync_memory() without return types