fix oom issues in export

Former-commit-id: b7ccc882a192aa1e25b1e5816f875ea304282412
2025-12-17 20:30:36 +08:00 · 2024-05-23 23:32:45 +08:00
parent a5a7d9ce95
commit 150c8022cc
2 changed files with 3 additions and 3 deletions
--- a/src/llamafactory/hparams/model_args.py
+++ b/src/llamafactory/hparams/model_args.py
@@ -145,7 +145,7 @@ class ModelArguments:
        default=1,
        metadata={"help": "The file shard size (in GB) of the exported model."},
    )
-    export_device: str = field(
+    export_device: Literal["cpu", "cuda"] = field(
        default="cpu",
        metadata={"help": "The device used in model export, use cuda to avoid addmm errors."},
    )