fix export on cpu device

Former-commit-id: cda2ff87272797a062c7addb1bf840ac46208dfd
This commit is contained in:
hiyouga 2024-03-04 17:35:09 +08:00
parent 0e58cd6422
commit a62d17d009

View File

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