return quantized indices as well

lucidrains · Oct 22, 2020 · 102b982 · 102b982
1 parent 2d5d324
commit 102b982
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/setup.py b/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',

diff --git a/vector_quantize_pytorch/vector_quantize_pytorch.py b/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