fix v6

JL-er · Mar 14, 2024 · c03cdbb · c03cdbb
1 parent 6d42247
commit c03cdbb
Show file tree

Hide file tree

Showing 8 changed files with 4 additions and 2 deletions.
diff --git a/.gitignore b/.gitignore
@@ -0,0 +1 @@
+__pycache__/
diff --git a/src/__pycache__/__init__.cpython-310.pyc b/src/__pycache__/__init__.cpython-310.pyc
diff --git a/src/__pycache__/binidx.cpython-310.pyc b/src/__pycache__/binidx.cpython-310.pyc
diff --git a/src/__pycache__/dataset.cpython-310.pyc b/src/__pycache__/dataset.cpython-310.pyc
diff --git a/src/__pycache__/model.cpython-310.pyc b/src/__pycache__/model.cpython-310.pyc
diff --git a/src/__pycache__/trainer.cpython-310.pyc b/src/__pycache__/trainer.cpython-310.pyc
diff --git a/src/__pycache__/utils.cpython-310.pyc b/src/__pycache__/utils.cpython-310.pyc
diff --git a/train.py b/train.py
@@ -257,8 +257,6 @@
     args.vocab_size = train_data.vocab_size
 
     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
@@ -267,6 +265,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():