Update vision_transformer.py

5 years ago · 28c0fa31fe
parent ee3d8050c8
commit 28c0fa31fe
1 changed files with 1 additions and 0 deletions
--- a/timm/models/vision_transformer.py
+++ b/timm/models/vision_transformer.py
@ -136,6 +136,7 @@ class Block(nn.Module):
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+        self.dropout = nn.Dropout(p=drop)

    def forward(self, x):
        residual = x.clone()