[REFACTOR][ToolAgent]

12 months ago · 537f81f932
parent 18b8a5d602
commit 537f81f932
1 changed files with 3 additions and 1 deletions
--- a/playground/agents/tool_agent.py
+++ b/playground/agents/tool_agent.py
@ -4,7 +4,9 @@ from swarms import ToolAgent

 # Load the pre-trained model and tokenizer
 model = AutoModelForCausalLM.from_pretrained(
-    "databricks/dolly-v2-12b"
+    "databricks/dolly-v2-12b",
+    load_in_4bit=True,
+    device_map="auto",
 )
 tokenizer = AutoTokenizer.from_pretrained("databricks/dolly-v2-12b")