fix export on cpu device

Former-commit-id: cda2ff8727
2026-01-02 12:10:34 +08:00 · 2024-03-04 17:35:09 +08:00
parent 0e58cd6422
commit a62d17d009
1 changed files with 2 additions and 1 deletions
--- a/src/llmtuner/model/patcher.py
+++ b/src/llmtuner/model/patcher.py
@@ -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(