bsq seems to work well

Author: lucidrains

examples/autoencoder_lfq.py

@@ -18,6 +18,8 @@
 codebook_size = 2 ** 8
 entropy_loss_weight = 0.02
 diversity_gamma = 1.
+spherical = True
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 class LFQAutoEncoder(nn.Module):
@@ -107,7 +109,8 @@ def iterate_dataset(data_loader):
 model = LFQAutoEncoder(
     codebook_size = codebook_size,
     entropy_loss_weight = entropy_loss_weight,
-    diversity_gamma = diversity_gamma
+    diversity_gamma = diversity_gamma,
+    spherical = spherical
 ).to(device)
 
 opt = torch.optim.AdamW(model.parameters(), lr=lr)

pyproject.toml

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

vector_quantize_pytorch/lookup_free_quantization.py

@@ -291,6 +291,10 @@ def forward(
         codebook_value = torch.ones_like(x) * self.codebook_scale
         quantized = torch.where(x > 0, codebook_value, -codebook_value)
 
+        # calculate indices
+
+        indices = reduce((quantized > 0).int() * self.mask.int(), 'b n c d -> b n c', 'sum')
+
         # maybe l2norm
 
         if self.spherical:
@@ -304,10 +308,6 @@ def forward(
         else:
             x = quantized
 
-        # calculate indices
-
-        indices = reduce((x > 0).int() * self.mask.int(), 'b n c d -> b n c', 'sum')
-
         # entropy aux loss
 
         if self.training: