0.1.12

Signed-off-by: ssbuild <[email protected]>
ssbuild · Jul 16, 2023 · ecf4fe5 · ecf4fe5
1 parent 81e9998
commit ecf4fe5
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/nlp/models/rwkv4/modeling_rwkv.py b/src/nlp/models/rwkv4/modeling_rwkv.py
@@ -684,6 +684,8 @@ def forward(self,
 
 
 class RwkvForCausalLM(RwkvPreTrainedModel):
+    _tied_weights_keys = ["head.weight"]
+
     def __init__(self, config: RwkvConfig):
         super().__init__(config)
         self.rwkv = RwkvModel(config)