when doing ema update in the presence of masked tokens, make sure to omit their contribution

Author: lucidrains

setup.py

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

vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -469,6 +469,10 @@ def forward(
             if self.affine_param:
                 flatten = (flatten - self.batch_mean) * (codebook_std / batch_std) + self.codebook_mean
 
+            if exists(mask):
+                mask = repeat(mask, 'b n -> h (b n)', h = flatten.shape[0])
+                embed_onehot[~mask] = 0.
+
             cluster_size = embed_onehot.sum(dim = 1)
 
             self.all_reduce_fn(cluster_size)
@@ -627,6 +631,10 @@ def forward(
             quantize = batched_embedding(embed_ind, embed)
 
         if self.training and self.ema_update:
+            if exists(mask):
+                mask = repeat(mask, 'b n -> h (b n)', h = flatten.shape[0])
+                embed_onehot[~mask] = 0.
+
             bins = embed_onehot.sum(dim = 1)
             self.all_reduce_fn(bins)