Rahul-8799 commited on
Commit
b1e1b89
·
verified ·
1 Parent(s): 16ab8a0

Update agents/project_manager_agent.py

Browse files
Files changed (1) hide show
  1. agents/project_manager_agent.py +1 -1
agents/project_manager_agent.py CHANGED
@@ -13,7 +13,7 @@ model = AutoModelForCausalLM.from_pretrained(
13
  def run(state):
14
  """Creates project plan based on product requirements."""
15
  prompt = state["input"]
16
- input_ids = tokenizer(prompt, return_tensors="pt").input_ids
17
  output_ids = model.generate(input_ids, max_new_tokens=512)
18
  output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
19
 
 
13
  def run(state):
14
  """Creates project plan based on product requirements."""
15
  prompt = state["input"]
16
+ input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
17
  output_ids = model.generate(input_ids, max_new_tokens=512)
18
  output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
19