code: def colo_memory_cap(size_in_GB):
      from colossalai.utils import (
          colo_device_memory_capacity,
          colo_set_process_memory_fraction,
          get_current_device,
      )
  
      cuda_capacity = colo_device_memory_capacity(get_current_device())
      if size_in_GB * (1024**3) < cuda_capacity:
          colo_set_process_memory_fraction(
              size_in_GB * (1024**3) / cuda_capacity
          )
          print(f"Limiting GPU memory usage to {size_in_GB} GB")
  
file_path: colossalai/examples/language/opt/opt_benchmark.py
imports: import time
  import torch
  import tqdm
  import transformers
  from args import parse_benchmark_args
  from transformers import AutoConfig, OPTForCausalLM
  from transformers.utils.versions import require_version
  import colossalai
  from colossalai.booster import Booster
  from colossalai.booster.plugin import GeminiPlugin, LowLevelZeroPlugin, TorchDDPPlugin
  from colossalai.cluster import DistCoordinator
  from colossalai.logging import disable_existing_loggers, get_dist_logger
  from colossalai.nn.optimizer import HybridAdam
label_desc: Use image node_id 3 for calling a global function with example usage: colo_memory_cap(size_in_GB) without return types