fix #3273

Former-commit-id: 3b20c89b342a068356ffc29c3724b645775c65db
2024-04-15 15:32:58 +08:00
parent 106a0104da
commit 9338f878a3
4 changed files with 18 additions and 8 deletions
--- a/src/llmtuner/hparams/model_args.py
+++ b/src/llmtuner/hparams/model_args.py
@@ -129,6 +129,10 @@ class ModelArguments:
        default=1,
        metadata={"help": "The file shard size (in GB) of the exported model."},
    )
+    export_device: str = field(
+        default="cpu",
+        metadata={"help": "The device used in model export."},
+    )
    export_quantization_bit: Optional[int] = field(
        default=None,
        metadata={"help": "The number of bits to quantize the exported model."},