code: def run_fn(
      init_method, model_fn, data_gen_fn, output_transform_fn
  ) -> Optional[str]:
      try:
          if init_method == "lazy":
              ctx = LazyInitContext()
          else:
              ctx = nullcontext()
          plugin = HybridParallelPlugin(
              tp_size=2, pp_size=2, num_microbatches=4, precision="bf16"
          )
          booster = Booster(plugin=plugin)
          with ctx:
              model = model_fn()
          optimizer = HybridAdam(model.parameters(), lr=1e-3)
          criterion = lambda x: x.mean()
          data = data_gen_fn()
  
          data = {
              k: v.to("cuda").repeat(4, 1)
              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
          )
  
          data_iter = iter([data])
  
          def _criterion(outputs, inputs):
              outputs = output_transform_fn(outputs)
              output_key = list(outputs.keys())[0]
              loss = criterion(outputs[output_key])
              return loss
  
          booster.execute_pipeline(
              data_iter,
              model,
              _criterion,
              optimizer,
              return_loss=True,
              return_outputs=False,
          )
          optimizer.step()
  
      except Exception as e:
          return repr(e)
  
imports: import copy
  from contextlib import nullcontext
  from typing import Optional
  import torch
  import torch.distributed
  from torch.testing import assert_close
  from torch.utils.data import Dataset
  import colossalai
  from colossalai.booster import Booster
  from colossalai.booster.plugin import HybridParallelPlugin
  from colossalai.fx import is_compatible_with_meta
  from colossalai.lazy.lazy_init import LazyInitContext
  from colossalai.nn.optimizer import HybridAdam
  from colossalai.testing import clear_cache_before_run, parameterize, rerun_if_address_is_in_use, spawn
  from colossalai.utils import get_current_device, set_seed
  from tests.kit.model_zoo import model_zoo
label_desc: Use image node_id 2 for calling a global function with example usage: run_fn(init_method, model_fn, data_gen_fn, output_transform_fn) and returns: loss, repr