def METHOD_NAME():
    MODEL_DICT = {tm.resnet18: [2100, 3000], tm.densenet121: [8100, 17000]}
    tracer = ColoTracer()
    for M, budgets in MODEL_DICT.items():
        for budget in budgets:
            model = M()
            graph = tracer.trace(model)
            gm = ColoGraphModule(model, graph, model.__class__.__name__)
            gm.graph._python_code = python_code_with_activation_checkpoint.__get__(graph)
            node_list = linearize(gm)
            gm = solver_rotor(gm, data=torch.rand(128, 3, 224, 224, device="meta"), mem_limit=budget * 1024**2)
            op_list = gm.__sequence__.list_operations()
            loss_op = next(op for op in op_list if isinstance(op, Loss))
            op_list = op_list[:op_list.index(loss_op)]
            in_ckpt = False
            ckpt_idx = 0
            for idx, op in enumerate(op_list):
                if in_ckpt:
                    if isinstance(op, ForwardNograd):
                        for n in node_list[idx]:
                            assert hasattr(n, "activation_checkpoint"), f"{n} is not annotated!"
                            assert n.activation_checkpoint == ckpt_idx, f"{n} ckpt_idx wrong, should be {ckpt_idx}!"
                        continue
                    if isinstance(op, ForwardEnable):
                        for n in node_list[idx]:
                            assert getattr(n, "activation_checkpoint", None) == None, f"{n} should not be annotated!"
                            in_ckpt = False
                        ckpt_idx += 1
                        continue
                    if isinstance(op, ForwardCheck):
                        ckpt_idx += 1
                        for n in node_list[idx]:
                            assert hasattr(n, "activation_checkpoint"), f"{n} is not annotated!"
                            assert n.activation_checkpoint == ckpt_idx, f"{n} ckpt_idx wrong, should be {ckpt_idx}!"
                        continue
                else:
                    if isinstance(op, ForwardCheck):
                        in_ckpt = True
                        for n in node_list[idx]:
                            assert hasattr(n, "activation_checkpoint"), f"{n} is not annotated!"
                            assert n.activation_checkpoint == ckpt_idx, f"{n} ckpt_idx wrong, should be {ckpt_idx}!"
            del model
            del gm
            del node_list