Add docstring

Signed-off-by: Jan Lasek <janek.lasek@gmail.com>
NVIDIA · Dec 23, 2024 · e24a4dd · e24a4dd
1 parent 0c1e9be
commit e24a4dd
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/nemo/collections/llm/recipes/run/executor.py b/nemo/collections/llm/recipes/run/executor.py
@@ -19,7 +19,15 @@
 
 @run.cli.factory
 def torchrun(devices: Optional[int] = None) -> run.Config[run.LocalExecutor]:
-    """Local executor using torchrun."""
+    """
+    Local executor using torchrun.
+
+    Args:
+        devices (Optional[int]): Number of devices to use. If None, it will use all available CUDA devices.
+
+    Returns:
+        run.Config[run.LocalExecutor]: Configuration for the local executor using torchrun.
+    """
     env_vars = {
         "TORCH_NCCL_AVOID_RECORD_STREAMS": "1",
     }