fix oom issues in export

2025-12-15 11:20:35 +08:00 · 2024-05-23 23:32:45 +08:00
parent 7134fb02bb
commit 67ebc7b388
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."},
    )