code: def test_nvme_adam(nvme_offload_fraction, nvme_offload_dir, adam_cls):
      model_builder, data_gen_fn, *_ = next(
          iter(model_zoo.get_sub_registry("custom_simple_net").values())
      )
      model = model_builder()
      torch_model = model_builder()
      move_some_params_to_cuda(model, torch_model)
      optimizer = adam_cls(
          model.parameters(),
          lr=0.1,
          nvme_offload_fraction=nvme_offload_fraction,
          nvme_offload_dir=nvme_offload_dir,
      )
      torch_optimizer = torch.optim.Adam(
          torch_model.parameters(), lr=0.1
      )
      with torch.no_grad():
          for p, torch_p in zip(
              model.parameters(), torch_model.parameters()
          ):
              torch_p.copy_(p)
              p.grad = torch.rand_like(p)
              torch_p.grad = p.grad
  
          for _ in range(3):
              optimizer.step()
              torch_optimizer.step()
              check_params_equal(model, torch_model)
  
imports: import torch
  from colossalai.nn.optimizer import CPUAdam, HybridAdam
  from colossalai.testing import clear_cache_before_run, parameterize
  from tests.kit.model_zoo import model_zoo
label_desc: Use image node_id 3 for calling a global function with example usage: test_nvme_adam(nvme_offload_fraction, nvme_offload_dir, adam_cls) without return types