docstr: Solve the checkpointing problem using Algorithm 3.
  
  Returns:
      graph (Graph): The optimized graph, should be a copy of the original graph.
docstr_tok: ["Solve","the","checkpointing","problem","using","Algorithm","3",".","Returns",":","graph","(","Graph",")",":","The","optimized","graph",",","should","be","a","copy","of","the","original","graph","."]
code: def solve(self) -> Graph:
      """Solve the checkpointing problem using Algorithm 3.
  
      Returns:
          graph (Graph): The optimized graph, should be a copy of the original graph.
      """
      checkpointable_op = [
          "call_module",
          "call_method",
          "call_function",
          "get_attr",
      ]
      ckpt = self.grid_search()
      for i, seg in enumerate(ckpt):
          for idx in range(*seg):
              nodes = self.node_list[idx]
              for n in nodes:
                  if n.op in checkpointable_op:
                      n.meta["activation_checkpoint"] = i
      return deepcopy(self.graph)
  
file_path: ColossalAI/colossalai/auto_parallel/checkpoint/ckpt_solver_chen.py
imports: import math
  from copy import deepcopy
  from typing import List, Set, Tuple
  from torch.fx import Graph, Node
  from colossalai.fx.profiler import calculate_fwd_in, calculate_fwd_tmp
  from .ckpt_solver_base import CheckpointSolverBase
label_desc: Use image node_id 2 for calling the CheckpointSolverChen obj's underlying member method code with example usage: obj.solve() and returns: deepcopy