code: def test_forward_easy(self) -> None:
      bsz = 3
      dim, size = 2, 9
      n = 33 * bsz
      memory_bank = MemoryBankModule(size=size)
  
      ptr = 0
      for i in range(0, n, bsz):
          output = torch.randn(2 * bsz, dim)
          output.requires_grad = True
          out0, out1 = output[:bsz], output[bsz:]
  
          _, curr_memory_bank = memory_bank(out1, update=True)
          next_memory_bank = memory_bank.bank.transpose(0, -1)
  
          curr_diff = out0.T - curr_memory_bank[:, ptr : ptr + bsz]
          next_diff = out1.T - next_memory_bank[:, ptr : ptr + bsz]
  
          # the current memory bank should not hold the batch yet
          self.assertGreater(curr_diff.norm(), 1e-5)
          # the "next" memory bank should hold the batch
          self.assertGreater(1e-5, next_diff.norm())
  
          ptr = (ptr + bsz) % size
  
file_path: lightly/tests/models/modules/test_memory_bank.py
imports: import re
  import unittest
  import pytest
  import torch
  from lightly.models.modules.memory_bank import MemoryBankModule
label_desc: Use image node_id 2 for calling the TestNTXentLoss obj's underlying member method code with example usage: obj.test_forward_easy() without return types