Update agents/software_architect_agent.py

agents/software_architect_agent.py

@@ -13,7 +13,7 @@ model = AutoModelForCausalLM.from_pretrained(
 def run(state):
     """Designs system architecture from project plan."""
     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)