pp

Signed-off-by: Mayank Mishra <[email protected]>
IBM · Oct 28, 2024 · 342ba39 · 342ba39
1 parent e419da4
commit 342ba39
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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