Put everything on the cuda device!

app.py

@@ -17,14 +17,19 @@ def chat_with_model(messages):
         yield messages + [{"role": "assistant", "content": "⚠️ No model loaded."}]
         return
 
-    current_model.to("cuda").half()
+    
 
     pad_id = current_tokenizer.pad_token_id
     if pad_id is None:
         pad_id = current_tokenizer.unk_token_id or 0
 
     prompt = format_prompt(messages)
-    inputs = current_tokenizer(prompt, return_tensors="pt").to(current_model.device)
+    device = torch.device("cuda")
+    current_model.to(device).half()
+
+    inputs = current_tokenizer(prompt, return_tensors="pt")
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+
 
     streamer = TextIteratorStreamer(current_tokenizer, skip_prompt=True, skip_special_tokens=False)