diff --git a/intermediate_source/pipelining_tutorial.rst b/intermediate_source/pipelining_tutorial.rst index 36738011a4..63170e6064 100644 --- a/intermediate_source/pipelining_tutorial.rst +++ b/intermediate_source/pipelining_tutorial.rst @@ -145,6 +145,7 @@ we are splitting before the before 4th transformer decoder layer, mirroring the we can retrieve a ``PipelineStage`` by calling ``build_stage`` after this splitting is done. .. code:: python + def tracer_model_split(model, example_input_microbatch) -> PipelineStage: pipe = pipeline( module=model,