Update agents/product_manager_agent.py

agents/product_manager_agent.py

@@ -13,9 +13,9 @@ model = AutoModelForCausalLM.from_pretrained(
 def run(state):
     """Generates structured product requirements from user input prompt."""
     prompt = state["input"]
-    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
     output_ids = model.generate(input_ids, max_new_tokens=512)
     output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
 
     new_log = state["chat_log"] + [{"role": "Product Manager", "content": output}]
-    return {"pm_output": output, "chat_log": new_log}
+    return {"pm_output": output, "chat_log": new_log}