def METHOD_NAME(
    self,
    model,
    graph,
    sharding_specs=None,
    tensor_shapes=None,
    device_topo=None,
    optimizer=None,
    nstages=None,
    **kwargs,
):
    nstages = device_topo.num_devices() if nstages is None else nstages
    return split_into_nstages_equal_size(model, graph, nstages)