quick fix for get_codes_from_indices

lucidrains · Mar 19, 2024 · 0f970c4 · 0f970c4
1 parent 6102e37
commit 0f970c4
Show file tree

Hide file tree

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

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