Update transformer.py

hehefan · Mar 30, 2022 · ea79fae · ea79fae
1 parent 03b9f01
commit ea79fae
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/modules-pytorch-1.8.1/transformer.py b/modules-pytorch-1.8.1/transformer.py
@@ -70,7 +70,7 @@ def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):
         self.layers = nn.ModuleList([])
         for _ in range(depth):
             self.layers.append(nn.ModuleList([
-                Residual(PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout))),
+                Residual(PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = 0.))),
                 Residual(PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout)))
             ]))
     def forward(self, x):