fix residual VQ in eval mode

Author: lucidrains

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'vector_quantize_pytorch',
   packages = find_packages(),
-  version = '0.3.8',
+  version = '0.3.9',
   license='MIT',
   description = 'Vector Quantization - Pytorch',
   author = 'Phil Wang',

vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -288,11 +288,11 @@ def forward(self, x):
 
         quantize, embed_ind = self._codebook(x)
 
-        commit_loss = 0.
-
         if self.training:
             commit_loss = F.mse_loss(quantize.detach(), x) * self.commitment
             quantize = x + (quantize - x).detach()
+        else:
+            commit_loss = torch.tensor([0.], device = x.device)
 
         quantize = self.project_out(quantize)