Initialize distributed runs if JAX_WORLD_SIZE > 1

Toni-SM · Jul 11, 2024 · 2db38e9 · 2db38e9
1 parent 828355d
commit 2db38e9
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/skrl/__init__.py b/skrl/__init__.py
@@ -124,9 +124,8 @@ def __init__(self) -> None:
                 # device
                 self._device = f"cuda:{self._local_rank}"
 
-                # TODO: find a better place for it
                 # set up distributed runs
-                if self._is_distributed and "jax" in sys.modules:
+                if self._is_distributed:
                     import jax
                     logger.info(f"Distributed (rank: {self._rank}, local rank: {self._local_rank}, world size: {self._world_size})")
                     jax.distributed.initialize(coordinator_address=self._coordinator_address,