From 342ba391fda568e132832ee5f759e50285c69301 Mon Sep 17 00:00:00 2001 From: Mayank Mishra Date: Mon, 28 Oct 2024 18:37:03 -0400 Subject: [PATCH] pp Signed-off-by: Mayank Mishra --- dolomite_engine/distributed/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dolomite_engine/distributed/__init__.py b/dolomite_engine/distributed/__init__.py index cba779a..40ecc9c 100644 --- a/dolomite_engine/distributed/__init__.py +++ b/dolomite_engine/distributed/__init__.py @@ -49,7 +49,7 @@ def wrap_model_container_for_distributed_training( model_container (ModelContainer): model container Returns: - tuple[model_container, _PipelineSchedule]: container of parallelized models and pipeline schedule + tuple[ModelContainer, _PipelineSchedule]: container of parallelized models and pipeline schedule """ stage = args.distributed_args.stage