Update my_model/KBVQA.py

my_model/KBVQA.py

@@ -66,8 +66,8 @@ class KBVQA():
 
     def load_fine_tuned_model(self):
 
-        self.kbvqa_model = AutoModelForCausalLM.from_pretrained(self.kbvqa_model_name, device_map="auto", quantization_config=self.bnb_config)
-        self.kbvqa_tokenizer = AutoTokenizer.from_pretrained(self.kbvqa_model_name, use_fast=self.use_fast, trust_remote_code=self.trust_remote, add_eos_token=self.add_eos_token)
+        self.kbvqa_model = AutoModelForCausalLM.from_pretrained(self.kbvqa_model_name, device_map="auto", low_cpu_mem_usage=True, quantization_config=self.bnb_config)
+        self.kbvqa_tokenizer = AutoTokenizer.from_pretrained(self.kbvqa_model_name, use_fast=self.use_fast, low_cpu_mem_usage=True, trust_remote_code=self.trust_remote, add_eos_token=self.add_eos_token)
 
 
     @property