Catch invalid device index

Toni-SM · Jun 22, 2024 · 5559069 · 5559069
1 parent 1ad585e
commit 5559069
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/skrl/envs/wrappers/jax/base.py b/skrl/envs/wrappers/jax/base.py
@@ -26,7 +26,7 @@ def __init__(self, env: Any) -> None:
                 device_type, device_index = f"{self._env.device}:0".split(':')[:2]
                 try:
                     self.device = jax.devices(device_type)[int(device_index)]
-                except RuntimeError:
+                except (RuntimeError, IndexError):
                     self.device = None
             else:
                 self.device = self._env.device
@@ -148,7 +148,7 @@ def __init__(self, env: Any) -> None:
                 device_type, device_index = f"{self._env.device}:0".split(':')[:2]
                 try:
                     self.device = jax.devices(device_type)[int(device_index)]
-                except RuntimeError:
+                except (RuntimeError, IndexError):
                     self.device = None
             else:
                 self.device = self._env.device