fix dora inference

Former-commit-id: 21b3597b0a05169afe51e1609b532787a65ca8ea
2024-03-05 11:51:41 +08:00
parent 02b838b9b0
commit 3553e301dd
1 changed files with 2 additions and 0 deletions
--- a/src/llmtuner/model/patcher.py
+++ b/src/llmtuner/model/patcher.py
@@ -287,6 +287,8 @@ def patch_config(
        init_kwargs["low_cpu_mem_usage"] = True
        if is_trainable:
            init_kwargs["device_map"] = {"": get_current_device()}
+        elif model_args.export_dir is None:
+            init_kwargs["device_map"] = "auto"


 def patch_model(