Fix non-used device in jax to torch

Farama-Foundation · Jul 2, 2024 · c1ec71d · c1ec71d
1 parent b18e6d1
commit c1ec71d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/gymnasium/wrappers/jax_to_torch.py b/gymnasium/wrappers/jax_to_torch.py
@@ -119,9 +119,9 @@ def _jax_iterable_to_torch(
     if hasattr(value, "_make"):
         # namedtuple - underline used to prevent potential name conflicts
         # noinspection PyProtectedMember
-        return type(value)._make(jax_to_torch(v) for v in value)
+        return type(value)._make(jax_to_torch(v, device) for v in value)
     else:
-        return type(value)(jax_to_torch(v) for v in value)
+        return type(value)(jax_to_torch(v, device) for v in value)
 
 
 class JaxToTorch(gym.Wrapper, gym.utils.RecordConstructorArgs):