From b0ab30cec4e0fd8c48dba7b0b70a8134b9161116 Mon Sep 17 00:00:00 2001
From: Gnome Ann <>
Date: Sun, 14 Nov 2021 18:24:51 -0500
Subject: [PATCH] Re-enable GPU-only generation option

---
 aiserver.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/aiserver.py b/aiserver.py
index 18fe309b..5a245602 100644
--- a/aiserver.py
+++ b/aiserver.py
@@ -359,8 +359,8 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]):
             print("(slower than GPU-only but uses less VRAM) or between multiple GPUs")
             print("(allowing you to use the combined VRAM of all your GPUs).")
             print("Currently only GPT-Neo and GPT-J models support this feature.")
-            print("{0}Use hybrid generation or CPU-only generation?:  (Default hybrid){1}".format(colors.CYAN, colors.END))
-            print(f"    1 - Hybrid generation\n    2 - CPU\n")
+            print("{0}Use hybrid generation, GPU-only generation or CPU-only generation?:  (Default hybrid){1}".format(colors.CYAN, colors.END))
+            print(f"    1 - Hybrid generation\n    2 - GPU\n    3 - CPU\n")
         else:
             print("    1 - GPU\n    2 - CPU\n")
         genselected = False
@@ -382,6 +382,10 @@ if(not vars.model in ["InferKit", "Colab", "OAI", "ReadOnly"]):
                     vars.usegpu = True
                     genselected = True
             elif(genselect.isnumeric() and int(genselect) == 2):
+                vars.breakmodel = False
+                vars.usegpu = True
+                genselected = True
+            elif(genselect.isnumeric() and int(genselect) == 3):
                 vars.breakmodel = False
                 vars.usegpu = False
                 genselected = True