Update main.py

main.py

@@ -5,6 +5,7 @@ from threading import Thread
 import time
 import torch
 import gradio as gr
+import spaces
 from concept_guidance.chat_template import DEFAULT_CHAT_TEMPLATE
 from concept_guidance.patching import patch_model, load_weights
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, TextIteratorStreamer, Conversation
@@ -12,8 +13,9 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, TextIter
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 # device = "cpu"
+device = "cuda"
 
 # comment in/out the models you want to use
 # RAM requirements: ~16GB x #models (+ ~4GB overhead)
@@ -86,6 +88,7 @@ def add_user_prompt(user_message, history):
     history.append([user_message, None])
     return history
 
+@spaces.GPU
 @torch.no_grad()
 def generate_completion(
     history,