From 102b982d4d7c7e4f34042d2d5904b1e26981d664 Mon Sep 17 00:00:00 2001
From: Phil Wang <lucidrains@gmail.com>
Date: Wed, 21 Oct 2020 19:26:36 -0700
Subject: [PATCH] return quantized indices as well

---
 setup.py                                           | 2 +-
 vector_quantize_pytorch/vector_quantize_pytorch.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/setup.py b/setup.py
index 6f7460c..94aeed2 100644
--- 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
index c86bd0d..d54b79b 100644
--- 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