class_docstr: Computes the scores of the model embeddings against the `positive'
  and `negative' samples from the Memory Bank.
  This class does *NOT* compute the actual loss, just the scores,
  i.e., inner products followed by normalizations/exponentiation etc.
class_docstr_tok: ["Computes","the","scores","of","the","model","embeddings","against","the","`","positive'","and","`","negative","'","samples","from","the","Memory","Bank",".","This","class","does","*","NOT","*","compute","the","actual","loss",",","just","the","scores",",","i.e.",",","inner","products","followed","by","normalizations\/exponentiation","etc","."]
code: def compute_partition_function(self, out):
      num_items = self.memory.size(0)
      with torch.no_grad():
          batch_mean = out.mean()
          # NOTE: this relies of "mean" computation being stable and deterministic
          # across all nodes. Could be replaced with smarter ways.
          if (
              torch.distributed.is_available()
              and torch.distributed.is_initialized()
          ):
              batch_mean_gathered = gather_from_all(batch_mean)
              all_batch_mean = (
                  batch_mean_gathered.mean().squeeze().item()
              )
          else:
              all_batch_mean = batch_mean.item()
      self.params[2] = all_batch_mean * num_items
      Z = self.params[2].clone().detach().item()
      rank = get_rank()
      logging.info(
          f"Rank: {rank}; Normalization constant Z is set to {Z}"
      )
  
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 3 for calling the NCEAverage obj's underlying member method code with example usage: obj.compute_partition_function(out) without return types