Merge pull request #271 from conacts/issue-260-cuda-error

Issue 260: fixing cuda not available
PromtEngineer · Jul 26, 2023 · 85a417d · 85a417d
2 parents 511afd1 + 36b50f3
commit 85a417d
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/ingest.py b/ingest.py
@@ -3,6 +3,7 @@
 from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
 
 import click
+import torch
 from langchain.docstore.document import Document
 from langchain.embeddings import HuggingFaceInstructEmbeddings
 from langchain.text_splitter import Language, RecursiveCharacterTextSplitter
@@ -89,7 +90,7 @@ def split_documents(documents: list[Document]) -> tuple[list[Document], list[Doc
 @click.command()
 @click.option(
     "--device_type",
-    default="cuda",
+    default="cuda" if torch.cuda.is_available() else "cpu",
     type=click.Choice(
         [
             "cpu",

diff --git a/run_localGPT.py b/run_localGPT.py
@@ -131,7 +131,7 @@ def load_model(device_type, model_id, model_basename=None):
 @click.command()
 @click.option(
     "--device_type",
-    default="cuda",
+    default="cuda" if torch.cuda.is_available() else "cpu",
     type=click.Choice(
         [
             "cpu",