make sure the codebook at least receive gradients from sync update

Author: lucidrains

setup.py

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

vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -798,11 +798,13 @@ def forward(
         quantize, embed_ind, distances = self._codebook(x, sample_codebook_temp = sample_codebook_temp)
 
         if self.training:
-            quantize = x + (quantize - x).detach()
+            maybe_sync_update = 0.
 
             if self.sync_update_v > 0.:
                 # (21) in https://minyoungg.github.io/vqtorch/assets/draft_050523.pdf
-                quantize = quantize + self.sync_update_v * (quantize - quantize.detach())
+                maybe_sync_update = self.sync_update_v * (quantize - quantize.detach())
+                print(maybe_sync_update)
+            quantize = x + (quantize - x).detach() + maybe_sync_update
 
         # function for calculating cross entropy loss to distance matrix
         # used for (1) naturalspeech2 training residual vq latents to be close to the correct codes and (2) cross-entropy based commitment loss