Hopefully fix CUDA again???

RedTachyon · Sep 9, 2023 · 9694479 · 9694479
1 parent 2118e01
commit 9694479
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 6 deletions.
diff --git a/coltra/__init__.py b/coltra/__init__.py
@@ -6,5 +6,5 @@
 from coltra.buffers import Action, Observation
 from coltra import utils
 
-__version__ = "0.1.8"
+__version__ = "0.1.9"
 VERSION = __version__
diff --git a/coltra/models/base_models.py b/coltra/models/base_models.py
@@ -297,19 +297,19 @@ def latent(self, x: Tensor, hidden_state: Tuple[Tensor, Tensor]) -> Tensor:
         return x
 
     def get_initial_state(
-        self, batch_size: int = 1, requires_grad: bool = True
+        self, batch_size: int = 1, requires_grad: bool = True, device: str = "cpu"
     ) -> Tuple[Tensor, Tensor]:
         return (
             torch.zeros(
                 batch_size,
                 self.lstm_hidden_size,
                 requires_grad=requires_grad,
-                device=self.device,
+                device=device,
             ),
             torch.zeros(
                 batch_size,
                 self.lstm_hidden_size,
                 requires_grad=requires_grad,
-                device=self.device,
+                device=device,
             ),
         )
diff --git a/coltra/models/lstm_models.py b/coltra/models/lstm_models.py
@@ -166,8 +166,8 @@ def get_initial_state(
         self, batch_size: int = 1, requires_grad: bool = True
     ) -> Tuple:
         return self.policy_network.get_initial_state(
-            batch_size, requires_grad
-        ), self.value_network.get_initial_state(batch_size, requires_grad)
+            batch_size, requires_grad, self.device
+        ), self.value_network.get_initial_state(batch_size, requires_grad, self.device)
 
 
 class FlattenLSTMModel(LSTMModel):