make mlp hidden_features defaults to 4*in_features

2023-04-13 11:08:21 +08:00 · 2023-04-13 11:08:21 +08:00 · 8c42415664
commit 8c42415664
parent 853ff72963
1 changed files with 1 additions and 1 deletions
--- a/flash_attn/modules/mlp.py
+++ b/flash_attn/modules/mlp.py
@ -17,7 +17,7 @@ class Mlp(nn.Module):
        factory_kwargs = {'device': device, 'dtype': dtype}
        super().__init__()
        out_features = out_features or in_features
-        hidden_features = hidden_features or in_features
+        hidden_features = hidden_features or in_features * 4
        self.return_residual = return_residual
        self.fc1 = nn.Linear(in_features, hidden_features, **factory_kwargs)
        self.activation = activation