code: def test_forward__dim_first(self) -> None:
      torch.manual_seed(0)
      memory_bank = MemoryBankModule(
          size=(5, 2), feature_dim_first=True
      )
      x0 = torch.randn(3, 2)
      out0, bank0 = memory_bank(x0, update=True)
      assert bank0.shape == (2, 5)
      x1 = torch.randn(3, 2)
      out1, bank1 = memory_bank(x1, update=True)
      assert bank1.shape == (2, 5)
      assert bank1[:, :3].tolist() == x0.T.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 5 for calling the TestMemoryBank obj's underlying member method code with example usage: obj.test_forward__dim_first() without return types