code: def test_adam():
      batch_size, seq_len = kt.bs_sl()
      hidden_dim = kt.hidden_dim
      cus_p = kt.rand((batch_size, seq_len, hidden_dim * 32))
      cus_out_p = kt.rand((batch_size, seq_len, hidden_dim * 32))
      cus_exp_avg = kt.rand((batch_size, seq_len, hidden_dim * 32))
      cus_exp_avg_sq = kt.rand((batch_size, seq_len, hidden_dim * 32))
      cus_grad = kt.rand((batch_size, seq_len, hidden_dim * 32))
  
      base_p = cus_p.detach().clone()
      base_out_p = cus_out_p.detach().clone()
      base_exp_avg = cus_exp_avg.detach().clone()
      base_exp_avg_sq = cus_exp_avg_sq.detach().clone()
      base_grad = cus_grad.detach().clone()
  
      print("total parameters {}".format(cus_p.numel()))
  
      def custom():
          adam_module.adam(
              cus_p,
              cus_out_p,
              cus_exp_avg,
              cus_exp_avg_sq,
              cus_grad,
              5e-4,
              0.9,
              0.98,
              1e-8,
              1,
              10,
              1,
              1,
              1e-4,
          )
          return [cus_p]
  
      def baseline():
          adam_module.apex_adam(
              base_p,
              base_out_p,
              base_exp_avg,
              base_exp_avg_sq,
              base_grad,
              5e-4,
              0.9,
              0.98,
              1e-8,
              1,
              10,
              1,
              1,
              1e-4,
          )
          return [base_p]
  
      return custom, baseline
  
imports: import random
  import torch
  from torch._C import ScriptModule, dtype
  from torch.nn import functional
  from lightseq.training.ops.pytorch.builder import KernelBuilder, AdamBuilder
  from tests.util import TestDecorator
  from torch_crf import CRF
label_desc: Use image node_id 16 for calling a global function with example usage: test_adam() and returns: custom,  baseline, list, list