code: def _gen_save_tensors_hooks_context(offload_input=True) -> str:
      """Generate customized saved_tensors_hooks
      Args:
          offload_input (bool, optional): whether we need offload input, if offload_input=False,
          we will use self.pack_hook_no_input instead. Defaults to True.
      Returns:
          str: generated context
      """
  
      if offload_input:
          context = "with torch.autograd.graph.saved_tensors_hooks(self.pack_hook_input, self.unpack_hook):\n"
      else:
          context = "with torch.autograd.graph.saved_tensors_hooks(self.pack_hook_no_input, self.unpack_hook):\n"
      return context
  
file_path: colossalai/colossalai/fx/codegen/activation_checkpoint_codegen.py
imports: from typing import Any, Dict, Iterable, List, Tuple
  import torch
  import colossalai
label_desc: Use image node_id 2 for calling a global function with example usage: _gen_save_tensors_hooks_context(offload_input) and returns: context