From 97ad42efe65e0d1027c65cda7192258a30924d55 Mon Sep 17 00:00:00 2001 From: KoboldAI Dev Date: Sun, 2 May 2021 23:34:33 -0400 Subject: [PATCH] Added device selection to transformers pipeline request to utilize GPU for inference --- aiserver.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aiserver.py b/aiserver.py index 4e84daeb..0e98a53e 100644 --- a/aiserver.py +++ b/aiserver.py @@ -108,7 +108,7 @@ if(vars.model != "InferKit"): print("{0}Initializing transformers, please wait...{1}".format(colors.HEADER, colors.ENDC)) from transformers import pipeline, GPT2Tokenizer - generator = pipeline('text-generation', model=vars.model) + generator = pipeline('text-generation', model=vars.model, device=0) tokenizer = GPT2Tokenizer.from_pretrained(vars.model) print("{0}OK! {1} pipeline created!{2}".format(colors.OKGREEN, vars.model, colors.ENDC)) else: