code: def trace_activation_checkpoint(self, enabled: bool):
      if enabled:
          orig_ckpt_func = torch.utils.checkpoint.CheckpointFunction
  
          class PatchedCheckpointFunction(torch.autograd.Function):
              @staticmethod
              def forward(ctx, run_function, preserve_rng_state, *args):
                  # signal that the current tracing occurs within activation checkpoint part
                  self.inside_torch_checkpoint_func = True
                  out = run_function(*args)
                  self.inside_torch_checkpoint_func = False
                  self.act_ckpt_region_count += 1
                  return out
  
              @staticmethod
              def backward(ctx: Any, *grad_outputs: Any) -> Any:
                  raise NotImplementedError(
                      "We do not implement the backward pass as we only trace the forward pass."
                  )
  
          # override the checkpoint function
          torch.utils.checkpoint.CheckpointFunction = (
              PatchedCheckpointFunction
          )
      yield
  
      if enabled:
          # recover the checkpoint function upon exit
          torch.utils.checkpoint.CheckpointFunction = orig_ckpt_func
  
imports: import functools
  import inspect
  import operator
  from contextlib import contextmanager
  from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
  import torch
  from torch.fx import Graph, Node, Proxy, Tracer
  from torch.utils._pytree import tree_map
  from colossalai.fx import ColoGraphModule, compatibility, is_compatible_with_meta
  from colossalai.fx.tracer._tracer_utils import is_element_in_list
  from colossalai.fx.tracer.bias_addition_patch import func_to_func_dict, method_to_func_dict, module_to_func_dict
  from colossalai.fx.tracer.registry import bias_addition_function, bias_addition_method, bias_addition_module, meta_patched_function, meta_patched_module
label_desc: Use image node_id 6 for calling the ColoTracer obj's underlying member method code with example usage: obj.trace_activation_checkpoint(enabled) and returns: out