address #145 again

Author: lucidrains

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "vector-quantize-pytorch"
-version = "1.14.45"
+version = "1.14.46"
 description = "Vector Quantization - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

vector_quantize_pytorch/lookup_free_quantization.py

@@ -198,7 +198,7 @@ def __init__(
         bits = ((all_codes[..., None].int() & self.mask) != 0).float()
         codebook = self.bits_to_codes(bits)
 
-        self.register_buffer('codebook', codebook, persistent = False)
+        self.register_buffer('codebook', codebook.float(), persistent = False)
 
     def bits_to_codes(self, bits):
         return bits * self.codebook_scale * 2 - self.codebook_scale
@@ -257,6 +257,7 @@ def forward(
         c - number of codebook dim
         """
 
+        orig_dtype = x.dtype
         x = x.float()
 
         is_img_or_video = x.ndim >= 4
@@ -313,7 +314,7 @@ def forward(
         # entropy aux loss
 
         if self.training:
-            codebook = self.codebook
+            codebook = self.codebook.float()
 
             codebook = self.maybe_l2norm(codebook)
 
@@ -403,6 +404,12 @@ def forward(
 
         aux_loss = entropy_aux_loss * self.entropy_loss_weight + commit_loss * self.commitment_loss_weight
 
+        # restore original dtype
+
+        x = x.type(orig_dtype)
+
+        # returns
+
         ret = Return(x, indices, aux_loss)
 
         if not return_loss_breakdown: