tweak device param

alan-turing-institute · Oct 10, 2024 · 8abfcf9 · 8abfcf9
1 parent 144b818
commit 8abfcf9
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/autoemulate/emulators/gaussian_process_torch.py b/autoemulate/emulators/gaussian_process_torch.py
@@ -59,7 +59,7 @@ def __init__(
         max_epochs=50,
         normalize_y=True,
         # misc
-        device=None,
+        device="cpu",
         random_state=None,
     ):
         self.mean_module = mean_module
@@ -167,11 +167,7 @@ def fit(self, X, y):
                 ),
             ],
             verbose=0,
-            device=self.device
-            if self.device is not None
-            else "cuda"
-            if torch.cuda.is_available()
-            else "cpu",
+            device=self.device,
         )
         self.model_.fit(X, y)
         self.is_fitted_ = True