code: def run_torch_amp():
      """
      In this test, we compare the torch amp and apex amp implemented in colossalai
      """
  
      torch.backends.cudnn.benchmark = False
      torch.backends.cudnn.deterministic = True
  
      # create layer
      test_models = ["torchvision_resnet18", "custom_simple_net"]
      for test_name in test_models:
          model_builder, data_gen_fn, *_ = next(
              iter(model_zoo.get_sub_registry(test_name).values())
          )
  
          # create model
          torch_amp_model = model_builder().cuda()
          apex_amp_model = copy.deepcopy(torch_amp_model)
  
          # create optimizer
          # we use SGD here, since the correctness of gradient clipping can't be tested with Adam
          torch_amp_optimizer = torch.optim.SGD(
              torch_amp_model.parameters(), lr=1e-3
          )
          apex_amp_optimizer = torch.optim.SGD(
              apex_amp_model.parameters(), lr=1e-3
          )
  
          # inject torch and apex amp
          torch_amp_config = dict(init_scale=128, enabled=True)
          (
              torch_amp_model,
              torch_amp_optimizer,
              _,
          ) = convert_to_torch_amp(
              torch_amp_model,
              torch_amp_optimizer,
              amp_config=torch_amp_config,
          )
          apex_amp_config = dict(opt_level="O1", loss_scale=128)
          apex_amp_model, apex_amp_optimizer = convert_to_apex_amp(
              apex_amp_model, apex_amp_optimizer, apex_amp_config
          )
  
          # create data
          data = data_gen_fn()
          data = {
              k: v.cuda() if isinstance(v, torch.Tensor) else v
              for k, v in data.items()
          }
  
          # forward pass
          torch_amp_output = torch_amp_model(**data)
          apex_amp_output = apex_amp_model(**data)
          assert_close_loose(torch_amp_output, apex_amp_output)
  
          for torch_amp_param, apex_amp_param in zip(
              torch_amp_model.parameters(), apex_amp_model.parameters()
          ):
              assert_close_loose(torch_amp_param, apex_amp_param)
  
          # backward
          # use sum() to get big gradient
          torch_amp_optimizer.backward(torch_amp_output.sum())
          apex_amp_optimizer.backward(apex_amp_output.sum())
  
          # check grad
          # In apex amp, grad is not scaled before backward, but torch amp does
          for torch_amp_param, apex_amp_param in zip(
              torch_amp_model.parameters(), apex_amp_model.parameters()
          ):
              assert_close_loose(
                  torch_amp_param.grad,
                  apex_amp_param.grad * apex_amp_config["loss_scale"],
              )
  
          # clip gradient
          apex_amp_optimizer.clip_grad_norm(
              model=apex_amp_model, max_norm=1.0
          )
          torch_amp_optimizer.clip_grad_norm(
              model=torch_amp_model, max_norm=1.0
          )
  
          # step
          torch_amp_optimizer.step()
          apex_amp_optimizer.step()
  
          # check updated param and grad
          for torch_amp_param, apex_amp_param in zip(
              torch_amp_model.parameters(), apex_amp_model.parameters()
          ):
              assert_close_loose(
                  torch_amp_param.grad, apex_amp_param.grad
              )
              assert_close_loose(torch_amp_param, apex_amp_param)
  
file_path: colossalai/tests/test_legacy/test_amp/test_torch_fp16.py
imports: import copy
  import pytest
  import torch
  import colossalai
  from colossalai.legacy.amp import convert_to_apex_amp, convert_to_torch_amp
  from colossalai.testing import assert_close_loose, clear_cache_before_run, rerun_if_address_is_in_use, spawn
  from tests.kit.model_zoo import model_zoo
label_desc: Use image node_id 1 for calling a global function with example usage: run_torch_amp() without return types