Merge branch 'fla-org:main' into main

fla/layers/rwkv6.py

@@ -191,7 +191,8 @@ def __init__(
         input_dim: int,
         output_dim: int,
         low_rank_dim: int,
-        bias: Optional[bool] = True
+        bias: Optional[bool] = True,
+        activation: Optional[str] = 'tanh'
     ):
         super().__init__()
 
@@ -200,9 +201,20 @@ def __init__(
         self.low_rank_dim = low_rank_dim
         self.bias = bias
 
+        if activation is None:
+            self.activation = nn.Identity()
+        elif activation == 'sigmoid':
+            self.activation = nn.Sigmoid()
+        elif activation == 'tanh':
+            self.activation = nn.Tanh()
+        elif activation == 'relu':
+            self.activation = nn.ReLU()
+        else:
+            raise ValueError(f"Not supported activation `{activation}`.")
+
         self.lora = nn.Sequential(
             nn.Linear(input_dim, low_rank_dim, bias=False),
-            nn.Tanh(),
+            self.activation,
             nn.Linear(low_rank_dim, output_dim, bias=bias)
         )