Update app.py

app.py

@@ -47,7 +47,7 @@ Assistant: girlfriend;mother;father;friend
 # setup
 torch.set_grad_enabled(False)
 model_name = "TheBloke/OpenHermes-2.5-Mistral-7B-GPTQ"
-pipe = pipeline("text-generation", model=model_name, device='cpu')
+pipe = pipeline("text-generation", model=model_name, device='cuda')
 generate_kwargs = {'max_new_tokens': 20}
 
 def past_kv_to_device(past_kv, device, dtype):
@@ -57,7 +57,7 @@ def detach_past_kv(past_kv):
     return tuple((k.cpu().detach().numpy().tolist(), v.cpu().detach().numpy().tolist()) for k, v in past_kv)
 
 
-# @spaces.GPU
+@spaces.GPU
 def set_past_key_values():
     model, tokenizer = pipe.model, pipe.tokenizer
     tokenized = tokenizer.encode(
@@ -90,7 +90,6 @@ def generate(text, past_key_values):
 if __name__ == "__main__":
     with torch.no_grad():
         past_key_values = set_past_key_values()
-        # pipe.model = pipe.model.cpu()
         demo = gr.Interface(
             partial(generate, past_key_values=past_key_values), 
             inputs="textbox", outputs="textbox"