quick fix for get_codes_from_indices

Author: lucidrains

setup.py

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

vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -826,8 +826,7 @@ def get_codes_from_indices(self, indices):
         is_multiheaded = codebook.ndim > 2
 
         if not is_multiheaded:
-            codes = codebook[indices]
-            return rearrange(codes, '... h d -> ... (h d)')
+            return codebook[indices]
 
         indices, ps = pack_one(indices, 'b * h')
         indices = rearrange(indices, 'b n h -> b h n')