cuda

app.py

@@ -26,7 +26,7 @@ def predict(history, max_length, top_p, temperature):
     messages = preprocess_messages(history)
     model_inputs = tokenizer.apply_chat_template(
         messages, add_generation_prompt=True, tokenize=True, return_tensors="pt", return_dict=True
-    )
+    ).to(model.device)
     streamer = TextIteratorStreamer(tokenizer, timeout=60, skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = {
         "input_ids": model_inputs["input_ids"],