Update main.py

main.py

@@ -100,7 +100,7 @@ def predict(message: str, history: list[str]) -> str:
         tokenize=False,
         add_generation_prompt=True
     )
-    model_inputs = tokenizer([text], return_tensors="pt")
+    model_inputs = tokenizer([text], return_tensors="pt").to("cuda:0")
 
     generate_kwargs = dict(
         model_inputs,