fix bug for post attention layer norms when using emb factorization

lucidrains · Nov 4, 2020 · 0444ff4 · 0444ff4
1 parent e86dffc
commit 0444ff4
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/setup.py b/setup.py
@@ -3,7 +3,7 @@
 setup(
   name = 'sinkhorn_transformer',
   packages = find_packages(exclude=['examples']),
-  version = '0.11.0',
+  version = '0.11.1',
   license='MIT',
   description = 'Sinkhorn Transformer - Sparse Sinkhorn Attention',
   author = 'Phil Wang',

diff --git a/sinkhorn_transformer/sinkhorn_transformer.py b/sinkhorn_transformer/sinkhorn_transformer.py
@@ -718,7 +718,7 @@ def __init__(self, num_tokens, dim, max_seq_len, depth, heads = 8, dim_head = No
         if emb_dim != dim:
             self.sinkhorn_transformer = ProjectInOut(self.sinkhorn_transformer, emb_dim, dim, project_out =(not return_embeddings))
 
-        self.norm = nn.LayerNorm(dim)
+        self.norm = nn.LayerNorm(emb_dim)
         self.to_logits = identity if return_embeddings else nn.Linear(emb_dim, num_tokens)
 
     def forward(self, x, **kwargs):