should be using l2norm of input for cosine sim

Author: lucidrains

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'vector_quantize_pytorch',
   packages = find_packages(),
-  version = '1.5.4',
+  version = '1.5.5',
   license='MIT',
   description = 'Vector Quantization - Pytorch',
   long_description_content_type = 'text/markdown',

vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -312,6 +312,9 @@ def forward(self, x):
 
         quantize = batched_embedding(embed_ind, self.embed)
 
+        if self.training:
+            quantize = x + (quantize - x).detach()
+
         if needs_codebook_dim:
             quantize, embed_ind = map(lambda t: rearrange(t, '1 ... -> ...'), (quantize, embed_ind))
 
@@ -456,6 +459,10 @@ def forward(self, x):
 
         quantize = batched_embedding(embed_ind, self.embed)
 
+        if self.training:
+            l2norm_x = l2norm(x)
+            quantize = l2norm_x + (quantize - l2norm_x).detach()
+
         if needs_codebook_dim:
             quantize, embed_ind = map(lambda t: rearrange(t, '1 ... -> ...'), (quantize, embed_ind))
 
@@ -587,9 +594,6 @@ def forward(
 
         quantize, embed_ind, distances = self._codebook(x)
 
-        if self.training:
-            quantize = x + (quantize - x).detach()
-
         if return_loss:
             if not is_multiheaded:
                 dist_einops_eq = '1 b n l -> b l n'