def METHOD_NAME(model, data_kwargs, num_stages, num_microbatches):
    tracer = ColoTracer()
    meta_args = {k: v.to('meta') for k, v in data_kwargs.items()}
    graph = tracer.trace(root=model, meta_args=meta_args)
    gm = torch.fx.GraphModule(model, graph, model.__class__.__name__)
    interp_meta_args = tuple([v.to('meta') for k, v in data_kwargs.items()])
    interp = MetaInfoProp(gm)
    interp.run(*interp_meta_args)
    #annotated_model = avgnode_split_pass(gm, num_stages)
    annotated_model = gpipe_dp_split_pass(gm, num_stages, num_microbatches, mode='block', block_limit=0.01)
    return annotated_model