code: def _benchmark_memory(model, inputs):
      with torch.no_grad():
          torch.cuda.reset_peak_memory_stats()
          now_mem = torch.cuda.memory_allocated() / 1024**2
          model(*inputs)
          new_max_mem = torch.cuda.max_memory_allocated() / 1024**2
      return new_max_mem - now_mem
  
file_path: colossalai/tests/test_autochunk/test_autochunk_alphafold/benchmark_autochunk_alphafold.py
imports: import time
  from typing import Any
  import torch
  import torch.fx
  import colossalai
  from colossalai.autochunk.autochunk_codegen import AUTOCHUNK_AVAILABLE
  from colossalai.fx.graph_module import ColoGraphModule
  from colossalai.fx.passes.meta_info_prop import MetaInfoProp
  from colossalai.testing import free_port
label_desc: Use image node_id 3 for calling a global function with example usage: _benchmark_memory(model, inputs) and returns: unknown