Fix #196 (#201)

MaartenGr · Jan 17, 2024 · 4916b2c · 4916b2c
1 parent f97900c
commit 4916b2c
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/keybert/_model.py b/keybert/_model.py
@@ -256,7 +256,9 @@ def extract_keywords(
         # Fine-tune keywords using an LLM
         if self.llm is not None:
             import torch
-            doc_embeddings = torch.from_numpy(doc_embeddings).float().to("cuda")
+            doc_embeddings = torch.from_numpy(doc_embeddings).float()
+            if torch.cuda.is_available():
+                doc_embeddings = doc_embeddings.to("cuda")
             if isinstance(all_keywords[0], tuple):
                 candidate_keywords = [[keyword[0] for keyword in all_keywords]]
             else: