Move torch-based import statements

Toni-SM · Oct 13, 2024 · 56439e0 · 56439e0
1 parent 155240b
commit 56439e0
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 8 deletions.
diff --git a/skrl/envs/wrappers/jax/isaacgym_envs.py b/skrl/envs/wrappers/jax/isaacgym_envs.py
@@ -12,15 +12,16 @@
     import torch.utils.dlpack as torch_dlpack
 except:
     pass  # TODO: show warning message
+else:
+    from skrl.utils.spaces.torch import (
+        convert_gym_space,
+        flatten_tensorized_space,
+        tensorize_space,
+        unflatten_tensorized_space
+    )
 
 from skrl import logger
 from skrl.envs.wrappers.jax.base import Wrapper
-from skrl.utils.spaces.torch import (
-    convert_gym_space,
-    flatten_tensorized_space,
-    tensorize_space,
-    unflatten_tensorized_space
-)
 
 
 # ML frameworks conversion utilities

diff --git a/skrl/envs/wrappers/jax/isaaclab_envs.py b/skrl/envs/wrappers/jax/isaaclab_envs.py
@@ -12,10 +12,11 @@
     import torch.utils.dlpack as torch_dlpack
 except:
     pass  # TODO: show warning message
+else:
+    from skrl.utils.spaces.torch import flatten_tensorized_space, tensorize_space, unflatten_tensorized_space
 
 from skrl import logger
 from skrl.envs.wrappers.jax.base import MultiAgentEnvWrapper, Wrapper
-from skrl.utils.spaces.torch import flatten_tensorized_space, tensorize_space, unflatten_tensorized_space
 
 
 # ML frameworks conversion utilities

diff --git a/skrl/envs/wrappers/jax/omniverse_isaacgym_envs.py b/skrl/envs/wrappers/jax/omniverse_isaacgym_envs.py
@@ -10,10 +10,11 @@
     import torch.utils.dlpack as torch_dlpack
 except:
     pass  # TODO: show warning message
+else:
+    from skrl.utils.spaces.torch import flatten_tensorized_space, tensorize_space, unflatten_tensorized_space
 
 from skrl import logger
 from skrl.envs.wrappers.jax.base import Wrapper
-from skrl.utils.spaces.torch import flatten_tensorized_space, tensorize_space, unflatten_tensorized_space
 
 
 # ML frameworks conversion utilities