code: def test():
      device_info = DeviceInfo(use_cuda=True, cuda_device=0)
      print(device_info)
      inp = torch.randn((64, 1, 28, 28), device=device_info.device)
      target = torch.ones(
          64, dtype=torch.long, device=device_info.device
      )
      loss_func = nn.CrossEntropyLoss()
      model = Model().to(device_info.device)
      my_adam = MyAdam(model.parameters())
      torch_adam = TorchAdam(model.parameters())
      loss = loss_func(model(inp), target)
      loss.backward()
      with monit.section("MyAdam warmup"):
          for i in range(100):
              my_adam.step()
      with monit.section("MyAdam"):
          for i in range(1000):
              my_adam.step()
      with monit.section("TorchAdam warmup"):
          for i in range(100):
              torch_adam.step()
      with monit.section("TorchAdam"):
          for i in range(1000):
              torch_adam.step()
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/optimizers/performance_test.py
imports: import torch
  import torch.nn
  from labml_helpers.device import DeviceInfo
  from torch.optim import Adam
  from labml import monit
  from labml_nn.optimizers.adam import Adam
  from labml_nn.optimizers.mnist_experiment import Model
label_desc: Use image node_id 1 for calling a global function with example usage: test() without return types