move dropout to be faithful to paper

Author: lucidrains

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__(

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',