From d17ce8461dccc7bc7914a0e882395fa3484040bc Mon Sep 17 00:00:00 2001 From: somebody Date: Wed, 12 Jul 2023 17:27:48 -0500 Subject: [PATCH] Use device_map="auto" --- modeling/inference_models/basic_hf/class.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modeling/inference_models/basic_hf/class.py b/modeling/inference_models/basic_hf/class.py index ecbc55cc..9d4b643b 100644 --- a/modeling/inference_models/basic_hf/class.py +++ b/modeling/inference_models/basic_hf/class.py @@ -55,7 +55,7 @@ class model_backend(HFInferenceModel): self.init_model_config() self.model = AutoModelForCausalLM.from_pretrained( - self.get_local_model_path(), low_cpu_mem_usage=True + self.get_local_model_path(), low_cpu_mem_usage=True, device_map="auto" ) if self.usegpu: