From a62d17d0092a339ef5783e7f5320b98c058b5ff7 Mon Sep 17 00:00:00 2001
From: hiyouga <hiyouga@buaa.edu.cn>
Date: Mon, 4 Mar 2024 17:35:09 +0800
Subject: [PATCH] fix export on cpu device

Former-commit-id: cda2ff87272797a062c7addb1bf840ac46208dfd
---
 src/llmtuner/model/patcher.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/llmtuner/model/patcher.py b/src/llmtuner/model/patcher.py
index aaedc1a8..e7ff0486 100644
--- 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(