Function split_into_stages 

pub fn split_into_stages(
    num_layers: usize,
    num_stages: usize,
) -> Vec<PipelineStage>

Split a model with num_layers layers into num_stages roughly equal stages.