class_docstr: Nearest Neighbour Memory Bank implementation
  
  This class implements a nearest neighbour memory bank as described in the
  NNCLR paper[0]. During the forward pass we return the nearest neighbour
  from the memory bank.
  
  [0] NNCLR, 2021, https://arxiv.org/abs/2104.14548
  
  Attributes:
      size:
          Size of the memory bank as (num_features, dim) tuple. If num_features is 0
          then the memory bank is disabled. Deprecated: If only a single integer is
          passed, it is interpreted as the number of features and the feature
          dimension is inferred from the first batch stored in the memory bank.
          Leaving out the feature dimension might lead to errors in distributed
          training.
  
  Examples:
      >>> model = NNCLR(backbone)
      >>> criterion = NTXentLoss(temperature=0.1)
      >>>
      >>> nn_replacer = NNmemoryBankModule(size=(2 ** 16, 128))
      >>>
      >>> # forward pass
      >>> (z0, p0), (z1, p1) = model(x0, x1)
      >>> z0 = nn_replacer(z0.detach(), update=False)
      >>> z1 = nn_replacer(z1.detach(), update=True)
      >>>
      >>> loss = 0.5 * (criterion(z0, p1) + criterion(z1, p0))
code: def __init__(self, size: Union[int, Sequence[int]] = 2**16):
      super(NNMemoryBankModule, self).__init__(size)
  
file_path: lightly/lightly/models/modules/nn_memory_bank.py
imports: from typing import Sequence, Union
  import torch
  from torch import Tensor
  from lightly.models.modules.memory_bank import MemoryBankModule
label_desc: Use image node_id 1 to create a new NNMemoryBankModule object from inherited base classes: MemoryBankModule with example: obj = NNMemoryBankModule(size)