Struct PipelineExecutor 

pub struct PipelineExecutor {
    pub stages: Vec<PipelineStage>,
    pub num_micro_batches: usize,
}

1F1B pipeline executor with per-micro-batch isolated tapes.

Fields

stages: Vec<PipelineStage>

num_micro_batches: usize

Implementations

impl PipelineExecutor

pub fn new(stages: Vec<PipelineStage>, num_micro_batches: usize) -> Self

pub fn run( &self, input: &Tensor, loss_fn: &(dyn Fn(&Tensor) -> Tensor + Send + Sync), ) -> Vec<GradientStore>

Run the pipeline. Returns per-stage gradient stores.