fix v6 lora (JL-er/RWKV-PEFT@c03cdbb)

josStorer · Mar 14, 2024 · c5077f4 · c5077f4
1 parent acf5d02
commit c5077f4
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/finetune/lora/v6/train.py b/finetune/lora/v6/train.py
@@ -314,8 +314,6 @@
 
  from src.model import RWKV, LORA_CONFIG, LoraLinear
 
- model = RWKV(args)
-
  if args.lora:
  assert args.lora_r > 0, "LoRA should have its `r` > 0"
  LORA_CONFIG["r"] = args.lora_r
@@ -324,6 +322,9 @@
  LORA_CONFIG["parts"] = set(str(args.lora_parts).split(","))
  enable_time_finetune = "time" in LORA_CONFIG["parts"]
  enable_ln_finetune = "ln" in LORA_CONFIG["parts"]
+ model = RWKV(args)
+
+ if args.lora:
  model.requires_grad_(False)
  for name, module in model.named_modules():