ros/LLaMA-Factory
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix export on cpu device

Browse Source 
Former-commit-id: e4722a9a627ea4e9a1341cc00a3108dd06a6b550
This commit is contained in:
hiyouga
2024-03-04 17:35:09 +08:00
parent f4f71f2797
commit 2dca53962e
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/llmtuner/model/patcher.py
View File

@@ -284,8 +284,9 @@ def patch_config(
init_kwargs["torch_dtype"] = model_args.compute_dtype init_kwargs["torch_dtype"] = model_args.compute_dtype
if not is_deepspeed_zero3_enabled(): if not is_deepspeed_zero3_enabled():
init_kwargs["device_map"] = {"": get_current_device()}
init_kwargs["low_cpu_mem_usage"] = True init_kwargs["low_cpu_mem_usage"] = True
if is_trainable:
init_kwargs["device_map"] = {"": get_current_device()}
def patch_model( def patch_model(