release the rotation trick for VQ from @cfifty https://arxiv.org/abs/2410.06424

Author: lucidrains

examples/autoencoder.py

@@ -88,6 +88,11 @@ def iterate_dataset(data_loader):
 
 print("baseline")
 torch.random.manual_seed(seed)
-model = SimpleVQAutoEncoder(codebook_size=num_codes).to(device)
+
+model = SimpleVQAutoEncoder(
+    codebook_size=num_codes,
+    rotation_trick=True
+).to(device)
+
 opt = torch.optim.AdamW(model.parameters(), lr=lr)
 train(model, train_dataset, train_iterations=train_iter)

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "vector-quantize-pytorch"
-version = "1.17.8"
+version = "1.18.0"
 description = "Vector Quantization - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }