in cosine sim, update with the l2norm of the smoothed embed_avg over smoothed cluster_size

Author: lucidrains

setup.py

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

vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -453,9 +453,7 @@ def forward(self, x):
             embed_normalized = self.embed_avg / rearrange(cluster_size, '... -> ... 1')
             embed_normalized = l2norm(embed_normalized)
 
-            self.embed.data.lerp_(embed_normalized, 1 - self.decay)
-            self.embed.data.copy_(l2norm(self.embed))
-
+            self.embed.data.copy_(l2norm(embed_normalized))
             self.expire_codes_(x)
 
         if needs_codebook_dim: