Update app.py

app.py

@@ -1,4 +1,5 @@
 import os
+import torch
 import threading
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
@@ -6,6 +7,8 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStream
 # Hugging Face token
 hf_token = os.environ["HUGGINGFACEHUB_API_TOKEN"]
 
+torch.set_num_threads(4)
+
 # Global model & tokenizer
 tokenizer = None
 model = None
@@ -36,7 +39,11 @@ def respond(message, max_tokens, temperature, top_p):
         eos_token_id=tokenizer.eos_token_id,
     )
 
-    thread = threading.Thread(target=model.generate, kwargs=generate_kwargs)
+    def run():
+        with torch.no_grad():
+            model.generate(**generate_kwargs)
+
+    thread = threading.Thread(target=run)
     thread.start()
 
     response = ""