Merge pull request #164 from ZhiyuanChen/patch-1

make mlp hidden_features defaults to 4*in_features
2023-04-12 23:21:12 -04:00 · 2023-04-12 23:21:12 -04:00 · 5cee071431
commit 5cee071431
parent 853ff72963 8c42415664
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