[Refactor] Use default device instead of CPU in losses

ghstack-source-id: 52a013a04a763bdb8c1c77a43a0984babe32bd77 Pull Request resolved: #2687
pytorch · Jan 10, 2025 · b7db5a5 · b7db5a5
1 parent 11305a7
commit b7db5a5
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
@@ -323,7 +323,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
@@ -306,7 +306,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
@@ -103,7 +103,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
@@ -195,7 +195,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
@@ -375,7 +375,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
@@ -309,7 +309,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
@@ -383,7 +383,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
@@ -1102,7 +1102,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):