Update app.py

app.py

@@ -54,9 +54,13 @@ def generate_response(model_id, conversation, user_message, max_length=512, temp
     # Create prompt
     prompt = "\n".join(conversation_history)
     
-    # Tokenize and generate
+    # Tokenize and move to GPU
     inputs = tokenizer(prompt, return_tensors="pt")
     
+    # Move inputs to the same device as the model
+    device = next(model.parameters()).device
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+    
     generation_start = time.time()
     with torch.no_grad():
         outputs = model.generate(