Coool2 commited on
Commit
a13e76f
·
1 Parent(s): 4e9b20a

Update agent.py

Browse files
Files changed (1) hide show
  1. agent.py +2 -5
agent.py CHANGED
@@ -209,12 +209,9 @@ def initialize_models(use_api_mode=False):
209
  # Embedding model
210
  embed_model = HuggingFaceEmbedding(
211
  model_name="llamaindex/vdr-2b-multi-v1",
212
- device="cpu",
213
  trust_remote_code=True,
214
- model_kwargs={
215
- "torch_dtype": "auto",
216
- "low_cpu_mem_usage": True
217
- }
218
  )
219
 
220
  return proj_llm, code_llm, embed_model
 
209
  # Embedding model
210
  embed_model = HuggingFaceEmbedding(
211
  model_name="llamaindex/vdr-2b-multi-v1",
212
+ device="cuda",
213
  trust_remote_code=True,
214
+ model_kwargs={"torch_dtype": "auto"},
 
 
 
215
  )
216
 
217
  return proj_llm, code_llm, embed_model