code: def run_fn(
      stage, model_fn, data_gen_fn, output_transform_fn
  ) -> Optional[str]:
      device = device_utils.get_current_device()
      try:
          plugin = LowLevelZeroPlugin(
              stage=stage, max_norm=1.0, initial_scale=2**5
          )
          booster = Booster(plugin=plugin)
          model = model_fn()
          optimizer = Adam(model.parameters(), lr=1e-3)
          criterion = lambda x: x.mean()
          data = data_gen_fn()
  
          data = {
              k: v.to(device)
              if torch.is_tensor(v) or "Tensor" in v.__class__.__name__
              else v
              for k, v in data.items()
          }
  
          model, optimizer, criterion, _, _ = booster.boost(
              model, optimizer, criterion
          )
  
          output = model(**data)
          output = output_transform_fn(output)
          output_key = list(output.keys())[0]
          loss = criterion(output[output_key])
  
          booster.backward(loss, optimizer)
          optimizer.step()
  
      except Exception as e:
          return repr(e)
  
file_path: ColossalAI/tests/test_booster/test_plugin/test_low_level_zero_plugin.py
imports: from typing import Optional
  import torch
  import torch.distributed
  from torch.optim import Adam
  import colossalai
  import colossalai.utils.device
  from colossalai.booster import Booster
  from colossalai.booster.plugin import LowLevelZeroPlugin
  from colossalai.testing import clear_cache_before_run, parameterize, rerun_if_address_is_in_use, spawn
  from tests.kit.model_zoo import COMMON_MODELS, IS_FAST_TEST, model_zoo
label_desc: Use image node_id 1 for calling a global function with example usage: run_fn(stage, model_fn, data_gen_fn, output_transform_fn) and returns: repr