address #116

Author: lucidrains

setup.py

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

vector_quantize_pytorch/finite_scalar_quantization.py

@@ -147,11 +147,6 @@ def forward(self, z: Tensor) -> Tensor:
         orig_dtype = z.dtype
         is_img_or_video = z.ndim >= 4
 
-        # make sure allowed dtype
-
-        if z.dtype not in self.allowed_dtypes:
-            z = z.float()
-
         # standardize image or video into (batch, seq, dimension)
 
         if is_img_or_video:
@@ -164,11 +159,23 @@ def forward(self, z: Tensor) -> Tensor:
 
         z = rearrange(z, 'b n (c d) -> b n c d', c = self.num_codebooks)
 
+        # make sure allowed dtype before quantizing
+
+        if z.dtype not in self.allowed_dtypes:
+            z = z.float()
+
         codes = self.quantize(z)
         indices = self.codes_to_indices(codes)
 
         codes = rearrange(codes, 'b n c d -> b n (c d)')
 
+        # cast codes back to original dtype
+
+        if codes.dtype != orig_dtype:
+            codes = codes.type(orig_dtype)
+
+        # project out
+
         out = self.project_out(codes)
 
         # reconstitute image or video dimensions
@@ -182,11 +189,6 @@ def forward(self, z: Tensor) -> Tensor:
         if not self.keep_num_codebooks_dim:
             indices = rearrange(indices, '... 1 -> ...')
 
-        # cast back to original dtype
-
-        if out.dtype != orig_dtype:
-            out = out.type(orig_dtype)
-
         # return quantized output and indices
 
         return out, indices