return quantized indices as well

Author: lucidrains

setup.py

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

vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -46,4 +46,4 @@ def forward(self, input):
 
         loss = F.mse_loss(quantize.detach(), input) * self.commitment
         quantize = input + (quantize - input).detach()
-        return quantize, loss
+        return quantize, embed_ind, loss