Specify different seeds in distributed runs for torch and jax

Toni-SM · Jul 11, 2024 · dc466d4 · dc466d4
1 parent 2db38e9
commit dc466d4
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/skrl/utils/__init__.py b/skrl/utils/__init__.py
@@ -74,8 +74,10 @@ def set_seed(seed: Optional[int] = None, deterministic: bool = False) -> int:
     seed = int(seed)
 
     # set different seeds in distributed runs
-    if config.torch.is_distributed or config.jax.is_distributed:
-        seed += max(config.torch.rank, config.jax.rank)
+    if config.torch.is_distributed:
+        seed += config.torch.rank
+    if config.jax.is_distributed:
+        seed += config.jax.rank
 
     logger.info(f"Seed: {seed}")