fix bug

2024-03-13 23:55:31 +08:00 · 2024-03-13 23:55:31 +08:00 · 8172530d54
parent 714d936dfb
commit 8172530d54
1 changed files with 2 additions and 0 deletions
--- a/src/llmtuner/train/utils.py
+++ b/src/llmtuner/train/utils.py
@ -164,6 +164,8 @@ def _create_galore_optimizer(

    if len(finetuning_args.galore_target) == 1 and finetuning_args.galore_target[0] == "all":
        galore_targets = find_all_linear_modules(model)
+    else:
+        galore_targets = finetuning_args.galore_target

    galore_params: List["torch.nn.Parameter"] = []
    for name, module in model.named_modules():