Update docstrings and docs

Toni-SM · Jan 24, 2025 · f830be9 · f830be9
1 parent e09ed4d
commit f830be9
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/docs/source/api/config/frameworks.rst b/docs/source/api/config/frameworks.rst
@@ -92,11 +92,12 @@ API
 
 .. py:data:: skrl.config.jax.device
     :type: jax.Device
-    :value: "cuda:${LOCAL_RANK}" | "cpu"
+    :value: "cuda:${JAX_LOCAL_RANK}" | "cpu"
 
     Default device.
 
-    The default device, unless specified, is ``cuda:0`` (or ``cuda:JAX_LOCAL_RANK`` in a distributed environment) if CUDA is available, ``cpu`` otherwise.
+    The default device, unless specified, is ``cuda:0`` if CUDA is available, ``cpu`` otherwise.
+    However, in a distributed environment, it is the device local to process with index ``JAX_RANK``.
 
 .. py:data:: skrl.config.jax.backend
     :type: str

diff --git a/skrl/__init__.py b/skrl/__init__.py
@@ -203,6 +203,10 @@ def parse_device(device: Union[str, "jax.Device", None]) -> "jax.Device":
 
                     This function supports the PyTorch-like ``"type:ordinal"`` string specification (e.g.: ``"cuda:0"``).
 
+                .. warning::
+
+                    This method returns (forces to use) the device local to process in a distributed environment.
+
                 :param device: Device specification. If the specified device is ``None`` or it cannot be resolved,
                                the default available device will be returned instead.
 
@@ -233,8 +237,8 @@ def parse_device(device: Union[str, "jax.Device", None]) -> "jax.Device":
             def device(self) -> "jax.Device":
                 """Default device.
 
-                The default device, unless specified, is ``cuda:0`` (or ``cuda:JAX_LOCAL_RANK`` in a distributed environment)
-                if CUDA is available, ``cpu`` otherwise.
+                The default device, unless specified, is ``cuda:0`` if CUDA is available, ``cpu`` otherwise.
+                However, in a distributed environment, it is the device local to process with index ``JAX_RANK``.
                 """
                 self._device = self.parse_device(self._device)
                 return self._device