move input ids and attmask to device

app.py

@@ -12,8 +12,13 @@ model = transformers.AutoModelForCausalLM.from_pretrained(model_name, device_map
 # Function to generate responses from the model
 def generate_response(input_text):
     inputs = tokenizer(input_text, return_tensors="pt")
-    outputs = model.generate(**inputs, max_new_tokens=512)
+    input_ids = inputs["input_ids"].to(model.device)
+    attention_mask = inputs["attention_mask"].to(model.device)
+
+    # Generate the output
+    outputs = model.generate(input_ids=input_ids, attention_mask=attention_mask, max_new_tokens=512)
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
     return response
 
 # Define the Gradio interface