[Refactor] Use default device instead of CPU in losses

ghstack-source-id: 8b98062c3ae88d8780ef7428fdfa07e305c790b9 Pull Request resolved: #2687
pytorch · Jan 16, 2025 · c3b9d1d · c3b9d1d
1 parent 256a700
commit c3b9d1d
Show file tree

Hide file tree

Showing 7 changed files with 8 additions and 8 deletions.
diff --git a/torchrl/objectives/cql.py b/torchrl/objectives/cql.py
@@ -324,7 +324,7 @@ def __init__(
         try:
             device = next(self.parameters()).device
         except AttributeError:
-            device = torch.device("cpu")
+            device = getattr(torch, "get_default_device", lambda: torch.device("cpu"))()
         self.register_buffer("alpha_init", torch.tensor(alpha_init, device=device))
         if bool(min_alpha) ^ bool(max_alpha):
             min_alpha = min_alpha if min_alpha else 0.0

diff --git a/torchrl/objectives/crossq.py b/torchrl/objectives/crossq.py
@@ -307,7 +307,7 @@ def __init__(
         try:
             device = next(self.parameters()).device
         except AttributeError:
-            device = torch.device("cpu")
+            device = getattr(torch, "get_default_device", lambda: torch.device("cpu"))()
         self.register_buffer("alpha_init", torch.tensor(alpha_init, device=device))
         if bool(min_alpha) ^ bool(max_alpha):
             min_alpha = min_alpha if min_alpha else 0.0

diff --git a/torchrl/objectives/decision_transformer.py b/torchrl/objectives/decision_transformer.py
@@ -104,7 +104,7 @@ def __init__(
         try:
             device = next(self.parameters()).device
         except AttributeError:
-            device = torch.device("cpu")
+            device = getattr(torch, "get_default_device", lambda: torch.device("cpu"))()
 
         self.register_buffer("alpha_init", torch.tensor(alpha_init, device=device))
         if bool(min_alpha) ^ bool(max_alpha):

diff --git a/torchrl/objectives/deprecated.py b/torchrl/objectives/deprecated.py
@@ -203,7 +203,7 @@ def __init__(
         try:
             device = next(self.parameters()).device
         except AttributeError:
-            device = torch.device("cpu")
+            device = getattr(torch, "get_default_device", lambda: torch.device("cpu"))()
 
         self.register_buffer("alpha_init", torch.as_tensor(alpha_init, device=device))
         self.register_buffer(

diff --git a/torchrl/objectives/ppo.py b/torchrl/objectives/ppo.py
@@ -393,7 +393,7 @@ def __init__(
         try:
             device = next(self.parameters()).device
         except (AttributeError, StopIteration):
-            device = torch.device("cpu")
+            device = getattr(torch, "get_default_device", lambda: torch.device("cpu"))()
 
         self.register_buffer("entropy_coef", torch.tensor(entropy_coef, device=device))
         if critic_coef is not None:

diff --git a/torchrl/objectives/redq.py b/torchrl/objectives/redq.py
@@ -319,7 +319,7 @@ def __init__(
         try:
             device = next(self.parameters()).device
         except AttributeError:
-            device = torch.device("cpu")
+            device = getattr(torch, "get_default_device", lambda: torch.device("cpu"))()
 
         self.register_buffer("alpha_init", torch.tensor(alpha_init, device=device))
         self.register_buffer(

diff --git a/torchrl/objectives/sac.py b/torchrl/objectives/sac.py
@@ -394,7 +394,7 @@ def __init__(
         try:
             device = next(self.parameters()).device
         except AttributeError:
-            device = torch.device("cpu")
+            device = getattr(torch, "get_default_device", lambda: torch.device("cpu"))()
         self.register_buffer("alpha_init", torch.tensor(alpha_init, device=device))
         if bool(min_alpha) ^ bool(max_alpha):
             min_alpha = min_alpha if min_alpha else 0.0
@@ -1121,7 +1121,7 @@ def __init__(
         try:
             device = next(self.parameters()).device
         except AttributeError:
-            device = torch.device("cpu")
+            device = getattr(torch, "get_default_device", lambda: torch.device("cpu"))()
 
         self.register_buffer("alpha_init", torch.tensor(alpha_init, device=device))
         if bool(min_alpha) ^ bool(max_alpha):