code: def run_fn(model_fn, data_gen_fn, output_transform_fn):
      plugin = TorchFSDPPlugin()
      booster = Booster(plugin=plugin)
      model = model_fn()
      optimizer = SGD(model.parameters(), lr=1e-3)
      criterion = lambda x: x.mean()
      data = data_gen_fn()
  
      data = {
          k: v.to("cuda")
          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
      )
  
      assert isinstance(model.module, FSDP)
      assert isinstance(optimizer, OptimizerWrapper)
  
      output = model(**data)
      output = output_transform_fn(output)
      output_key = list(output.keys())[0]
      loss = criterion(output[output_key])
  
      booster.backward(loss, optimizer)
      optimizer.clip_grad_by_norm(1.0)
      optimizer.step()
  
      del model
      del optimizer
      del criterion
      del booster
      del plugin
  
file_path: colossalai/tests/test_booster/test_plugin/test_torch_fsdp_plugin.py
imports: import pytest
  import torch
  from packaging import version
  from torch.optim import SGD
  import colossalai
  from colossalai.booster import Booster
  from colossalai.interface import OptimizerWrapper
  from colossalai.testing import clear_cache_before_run, 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(model_fn, data_gen_fn, output_transform_fn) without return types