code: def _main():
      device = "cpu"
      model_size = 1 * 1024**3
      group_size = [model_size, 274432]
      param = [
          torch.nn.Parameter(torch.ones(size, device=device))
          for size in group_size
      ]
      torch_time = _test_perf(param, torch.optim.Adagrad)
      ds_time = _test_perf(param, DeepSpeedCPUAdagrad)
      print(f"Step time: {torch_time=} {ds_time=}")
  
imports: import torch
  from deepspeed.ops.adagrad import DeepSpeedCPUAdagrad
  import time
label_desc: Use image node_id 2 for calling a global function with example usage: _main() without return types