Merge pull request #164 from ZhiyuanChen/patch-1

make mlp hidden_features defaults to 4*in_features

Merge pull request #164 from ZhiyuanChen/patch-1
make mlp hidden_features defaults to 4*in_features
5cee0714 · Tri Dao · GitHub · 853ff729 · 8c424156 · 5cee0714
Unverified Commit 5cee0714 authored Apr 12, 2023 by Tri Dao Committed by GitHub Apr 12, 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