code: def memory_optimize(
      model: torch.nn.Module,
      inps: Dict[str, torch.Tensor],
      memory_budget: float = -1.0,
      solver_name: str = "asyn",
  ):
      model = model.cpu().half()
      tracer = ColoTracer()
      assert is_compatible_with_meta()
      wrap_fn = (
          lambda x: x.to("meta") if isinstance(x, torch.Tensor) else x
      )
      meta_args = tree_map(wrap_fn, inps)
      graph = tracer.trace(model, meta_args=meta_args)
      gm = GraphModule(model, graph, model.__class__.__name__)
      interp = MetaInfoProp(gm)
      interp.propagate(*meta_args.values())
  
      region_manager = RegionManager(
          graph, solver_name=solver_name, memory_budget=memory_budget
      )
      region_manager._build_regions()
      GlobalRuntimeInfo().region_list = region_manager.region_list
  
      act_peak_mem = (
          compute_act_peak_mem(region_manager.region_list) / 1024**2
      )
      max_param_mem = (
          compute_max_param_mem(region_manager.region_list) / 1024**2
      )
      total_param_mem = (
          compute_total_param_mem(region_manager.region_list)
          / 1024**2
      )
      print(
          f"act_peak_mem={act_peak_mem:.3f} MB | max_param_mem={max_param_mem:.3f} MB | total_param_mem={total_param_mem:.3f}"
      )
  
      if solver_name == "syn":
          gm = runtime_syn_offload_apply_pass(
              gm, region_manager.region_list
          )
      elif solver_name == "asyn":
          gm = runtime_asyn_offload_apply_pass(
              gm, region_manager.region_list
          )
      else:
          raise TypeError(f"Unknown solver name {solver_name}!")
  
      gm.recompile()
      optimized_model = BaseOffloadModule(
          gm, region_manager, solver_name == "syn"
      )
      return optimized_model
  
file_path: ColossalAI/colossalai/auto_parallel/offload/mem_optimize.py
imports: from typing import Dict
  import torch
  import torch.fx
  from torch.fx import GraphModule
  from torch.utils._pytree import tree_map
  from colossalai.fx import ColoTracer, is_compatible_with_meta
  from colossalai.fx.passes.meta_info_prop import MetaInfoProp
  from .base_offload_module import BaseOffloadModule
  from .region_manager import RegionManager
  from .runtime import runtime_asyn_offload_apply_pass, runtime_syn_offload_apply_pass
  from .util import GlobalRuntimeInfo, compute_act_peak_mem, compute_max_param_mem, compute_total_param_mem
label_desc: Use image node_id 1 for calling a global function with example usage: memory_optimize(model, inps, memory_budget, solver_name) and returns: optimized_model