PipelineStage

class tensorplay.distributed.pipelining.PipelineStage(submodule: Any, stage_index: int, num_stages: int, device: Any = None, input_args: tuple[Any, ...] | None = None, output_args: Any = None, output_grads: Any = None, input_grads: Any = None, group: Any = None, dw_builder: Callable[[], Callable[[...], None]] | None = None, get_mesh: Any = None)[source]