Fix device arg for sentence transformers (#124)

deepset-ai · May 28, 2020 · 46a065e · 46a065e
1 parent df554fc
commit 46a065e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/haystack/retriever/elasticsearch.py b/haystack/retriever/elasticsearch.py
@@ -84,7 +84,7 @@ def __init__(
             # pretrained embedding models coming from: https://github.com/UKPLab/sentence-transformers#pretrained-models
             # e.g. 'roberta-base-nli-stsb-mean-tokens'
             if gpu:
-                device = "gpu"
+                device = "cuda"
             else:
                 device = "cpu"
             self.embedding_model = SentenceTransformer(embedding_model, device=device)