set requires_grad depending on whether training or not

Author: lucidrains

setup.py

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

vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -311,7 +311,7 @@ def forward(self, x):
         if self.training:
             quantize = x + (quantize - x).detach()
 
-        loss = torch.tensor([0.], device = device)
+        loss = torch.tensor([0.], device = device, requires_grad = self.training)
 
         if self.training:
             if self.commitment > 0: