init self.embed_avg correctly during kmeans init

Author: lucidrains

setup.py

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

vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -332,8 +332,10 @@ def init_embed_(self, data, mask = None):
             all_reduce_fn = self.kmeans_all_reduce_fn
         )
 
+        embed_sum = embed * rearrange(cluster_size, '... -> ... 1')
+
         self.embed.data.copy_(embed)
-        self.embed_avg.data.copy_(embed.clone())
+        self.embed_avg.data.copy_(embed_sum)
         self.cluster_size.data.copy_(cluster_size)
         self.initted.data.copy_(torch.Tensor([True]))
 
@@ -580,8 +582,10 @@ def init_embed_(self, data, mask = None):
             all_reduce_fn = self.kmeans_all_reduce_fn
         )
 
+        embed_sum = embed * rearrange(cluster_size, '... -> ... 1')
+
         self.embed.data.copy_(embed)
-        self.embed_avg.data.copy_(embed.clone())
+        self.embed_avg.data.copy_(embed_sum)
         self.cluster_size.data.copy_(cluster_size)
         self.initted.data.copy_(torch.Tensor([True]))