code: def test_cpu_adam_opt(dtype, model_size):
      if ("amd" in pytest.cpu_vendor) and (dtype == torch.half):
          pytest.skip(
              "cpu-adam with half precision not supported on AMD CPUs"
          )
  
      from deepspeed.ops.adam import DeepSpeedCPUAdam
  
      device = "cpu"
      rng_state = torch.get_rng_state()
      param = torch.nn.Parameter(
          torch.randn(model_size, device=device).to(dtype)
      )
      torch.set_rng_state(rng_state)
      param1_data = torch.randn(model_size, device=device)
      param1 = torch.nn.Parameter(param1_data)
      torch.set_rng_state(rng_state)
      param2_data = (
          torch.randn(model_size, device=device).to(dtype).cuda()
      )
      param2 = torch.nn.Parameter(param2_data)
  
      optimizer1 = torch.optim.AdamW([param1])
      optimizer2 = FusedAdam([param2])
      optimizer = DeepSpeedCPUAdam([param])
  
      for i in range(10):
          rng_state = torch.get_rng_state()
          param.grad = torch.randn(model_size, device=device).to(dtype)
          torch.set_rng_state(rng_state)
          param1.grad = torch.randn(model_size, device=device)
          torch.set_rng_state(rng_state)
          param2.grad = (
              torch.randn(model_size, device=device).to(dtype).cuda()
          )
  
          optimizer.step()
          optimizer2.step()
          optimizer1.step()
      tolerance = param1.float().norm().detach().numpy() * 1e-2
      check_equal(
          param.float().norm(),
          param1.float().norm(),
          atol=tolerance,
          verbose=True,
      )
      check_equal(
          param.float().norm(),
          param2.float().cpu().norm(),
          atol=tolerance,
          verbose=True,
      )
  
file_path: H2O/h2o_flexgen/benchmark/third_party/DeepSpeed/tests/unit/ops/adam/test_cpu_adam.py
imports: import torch
  import numpy
  import pytest
  from cpuinfo import get_cpu_info
  import deepspeed
  from deepspeed.ops.adam import FusedAdam
  from deepspeed.ops.op_builder import CPUAdamBuilder
label_desc: Use image node_id 2 for calling a global function with example usage: test_cpu_adam_opt(dtype, model_size) without return types