code: def test_forward(self) -> None:
      torch.manual_seed(0)
      memory_bank = MemoryBankModule(
          size=(5, 2), feature_dim_first=False
      )
      x0 = torch.randn(3, 2)
      out0, bank0 = memory_bank(x0, update=True)
      # Verify that output is same as input.
      assert out0.tolist() == x0.tolist()
      # Verify that memory bank was initialized and has correct shape.
      assert bank0.shape == (5, 2)
      assert memory_bank.bank.shape == (5, 2)
      # Verify that output bank does not contain features from x0.
      assert bank0[:3].tolist() != x0.tolist()
      # Verify that memory bank was updated.
      assert memory_bank.bank[:3].tolist() == x0.tolist()
  
      x1 = torch.randn(3, 2)
      out1, bank1 = memory_bank(x1, update=True)
      # Verify that output is same as input.
      assert out1.tolist() == x1.tolist()
      # Verify that output bank contains features from x0.
      assert bank1[:3].tolist() == x0.tolist()
      # Verify that output bank does not contain features from x1.
      assert bank1[3:].tolist() != x1[:2].tolist()
      # Verify that memory bank was updated.
      assert memory_bank.bank[:3].tolist() == x0.tolist()
      assert memory_bank.bank[3:].tolist() == x1[:2].tolist()
  
      # At this point the memory bank is full.
      # Adding more features will start overwriting the bank from the beginning.
  
      x2 = torch.randn(3, 2)
      out2, bank2 = memory_bank(x2, update=True)
      # Verify that output is same as input.
      assert out2.tolist() == x2.tolist()
      # Verify that output bank contains features from x0 and x1.
      assert bank2[:3].tolist() == x0.tolist()
      assert bank2[3:].tolist() == x1[:2].tolist()
      # Verify that memory bank is overwritten.
      assert memory_bank.bank[:3].tolist() == x2.tolist()
  
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 3 for calling the TestMemoryBank obj's underlying member method code with example usage: obj.test_forward() without return types