Limited CTMs to only use CPU at all times

x-tabdeveloping · Sep 10, 2024 · 1e74a44 · 1e74a44
1 parent 21fd848
commit 1e74a44
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/turftopic/models/ctm.py b/turftopic/models/ctm.py
@@ -213,9 +213,7 @@ def fit(
             seed = self.random_state or random.randint(0, 10_000)
             torch.manual_seed(seed)
             pyro.set_rng_seed(seed)
-            device = torch.device(
-                "cuda:0" if torch.cuda.is_available() else "cpu"
-            )
+            device = torch.device("cpu")
             pyro.clear_param_store()
             contextualized_size = embeddings.shape[1]
             if self.combined: