[BugFix] Fix get_default_device calls in older PT versions

ghstack-source-id: fd3a739
Pull Request resolved: #2586

Author: Vincent Moens

benchmarks/test_objectives_benchmarks.py

@@ -52,7 +52,7 @@
 
 @pytest.fixture(scope="module", autouse=True)
 def set_default_device():
-    cur_device = torch.get_default_device()
+    cur_device = getattr(torch, "get_default_device", lambda: torch.device("cpu"))()
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
     torch.set_default_device(device)
     yield

torchrl/envs/libs/vmas.py

@@ -803,7 +803,9 @@ def _build_env(
                 num_envs=num_envs,
                 device=self.device
                 if self.device is not None
-                else torch.get_default_device(),
+                else getattr(
+                    torch, "get_default_device", lambda: torch.device("cpu")
+                )(),
                 continuous_actions=continuous_actions,
                 max_steps=max_steps,
                 seed=seed,

torchrl/objectives/utils.py

@@ -596,7 +596,7 @@ def _get_default_device(net):
     for p in net.parameters():
         return p.device
     else:
-        return torch.get_default_device()
+        return getattr(torch, "get_default_device", lambda: torch.device("cpu"))()
 
 
 def group_optimizers(*optimizers: torch.optim.Optimizer) -> torch.optim.Optimizer:

torchrl/objectives/value/advantages.py

@@ -219,7 +219,7 @@ def __init__(
     ):
         super().__init__()
         if device is None:
-            device = torch.get_default_device()
+            device = getattr(torch, "get_default_device", lambda: torch.device("cpu"))()
         # this is saved for tracking only and should not be used to cast anything else than buffers during
         # init.
         self._device = device