Update merge_lora_params.py (#8514)

* Update merge_lora_params.py * Update merge_lora_params.py * Update merge_lora_params.py
PaddlePaddle · May 31, 2024 · 60c6615 · 60c6615
1 parent 92aa9f1
commit 60c6615
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/llm/merge_lora_params.py b/llm/merge_lora_params.py
@@ -125,7 +125,7 @@ def merge():
         model = AutoModelForCausalLM.from_pretrained(
             lora_config.base_model_name_or_path,
             config=config,
-            low_cpu_mem_usage=True,
+            low_cpu_mem_usage=args.low_gpu_mem,
         )
         lora_config.merge_weights = True
         model = LoRAModel.from_pretrained(model=model, lora_path=args.lora_path, lora_config=lora_config)