move dropout to be faithful to paper

lucidrains · Jan 4, 2021 · 86b8316 · 86b8316
1 parent 48954bb
commit 86b8316
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/conformer/conformer.py b/conformer/conformer.py
@@ -118,11 +118,11 @@ def forward(self, x, context = None, mask = None, context_mask = None):
             dots.masked_fill_(~mask, mask_value)
 
         attn = dots.softmax(dim = -1)
-        attn = self.dropout(attn)
 
         out = einsum('b h i j, b h j d -> b h i d', attn, v)
         out = rearrange(out, 'b h n d -> b n (h d)')
-        return self.to_out(out)
+        out = self.to_out(out)
+        return self.dropout(out)
 
 class FeedForward(nn.Module):
     def __init__(

diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'conformer',
   packages = find_packages(),
-  version = '0.2.2',
+  version = '0.2.3',
   license='MIT',
   description = 'The convolutional module from the Conformer paper',
   author = 'Phil Wang',