Update app.py

app.py

@@ -6,7 +6,7 @@ from transformers import StoppingCriteria, StoppingCriteriaList, TextIteratorStr
 from threading import Thread
 
 
-torch.set_default_device("cuda")
+
 
 # Loading the tokenizer and model from Hugging Face's model hub.
 tokenizer = AutoTokenizer.from_pretrained(
@@ -33,6 +33,8 @@ class StopOnTokens(StoppingCriteria):
 # Function to generate model predictions.
 @spaces.GPU
 def predict(message, history):
+    torch.set_default_device("cuda")
+    
     history_transformer_format = history + [[message, ""]]
     stop = StopOnTokens()