Added cuda testing

MushroomRL · Dec 4, 2023 · 9c24834 · 9c24834
1 parent c7e133e
commit 9c24834
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/tests/environments/test_vectorized_envs.py → tests/core/test_vectorized_envs.py b/tests/environments/test_vectorized_envs.py → tests/core/test_vectorized_envs.py
@@ -4,6 +4,7 @@
 from mushroom_rl.core import Agent, VectorCore, VectorizedEnvironment, MDPInfo
 from mushroom_rl.rl_utils import Box
 from mushroom_rl.policy import Policy
+from mushroom_rl.utils import TorchUtils
 
 
 class DummyPolicy(Policy):
@@ -94,3 +95,10 @@ def test_vectorized_env_():
     run_exp(env_backend='torch', agent_backend='numpy')
     run_exp(env_backend='numpy', agent_backend='torch')
     run_exp(env_backend='numpy', agent_backend='numpy')
+
+    if torch.cuda.is_available():
+        print('Testing also cuda')
+        TorchUtils.set_default_device('cuda')
+        run_exp(env_backend='torch', agent_backend='torch')
+        run_exp(env_backend='torch', agent_backend='numpy')
+