From 537f81f93272cd6e3b57ffb72e7b0b083807efde Mon Sep 17 00:00:00 2001
From: Kye <kye@apacmediasolutions.com>
Date: Wed, 31 Jan 2024 21:46:09 -0500
Subject: [PATCH] [REFACTOR][ToolAgent]

---
 playground/agents/tool_agent.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/playground/agents/tool_agent.py b/playground/agents/tool_agent.py
index 71d10bb8..0a95f42c 100644
--- 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")