Skip to content

Commit

Permalink
release the rotation trick for VQ from @cfifty https://arxiv.org/abs/…
Browse files Browse the repository at this point in the history
  • Loading branch information
lucidrains committed Oct 11, 2024
1 parent d3dd5e5 commit 089011b
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 2 deletions.
7 changes: 6 additions & 1 deletion examples/autoencoder.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
@@ -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]" }
Expand Down

0 comments on commit 089011b

Please sign in to comment.