make sure get_codes_from_indices work for grouped rvq

Author: lucidrains

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'vector_quantize_pytorch',
   packages = find_packages(),
-  version = '1.4.0',
+  version = '1.4.1',
   license='MIT',
   description = 'Vector Quantization - Pytorch',
   long_description_content_type = 'text/markdown',

vector_quantize_pytorch/residual_vq.py

@@ -220,6 +220,10 @@ def __init__(
     def codebooks(self):
         return torch.stack(tuple(rvq.codebooks for rvq in self.rvqs))
 
+    def get_codes_from_indices(self, indices):
+        codes = tuple(rvq.get_codes_from_indices(chunk_indices) for rvq, chunk_indices in zip(self.rvqs, indices))
+        return torch.stack(codes)
+
     def forward(
         self,
         x,