make mlp hidden_features defaults to 4*in_features

8c424156 · Zhiyuan Chen · GitHub · 853ff729 · 8c424156
Unverified Commit 8c424156 authored Apr 13, 2023 by Zhiyuan Chen Committed by GitHub Apr 13, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

flash_attn/modules/mlp.py flash_attn/modules/mlp.py +1 -1

No files found.
--- 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