quick fix

Author: lucidrains

pyproject.toml

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

tests/test_beam.py

@@ -51,13 +51,15 @@ def test_beam_search():
         dim = 256,
         num_quantizers = 8,      # specify number of quantizers
         codebook_size = 1024,    # codebook size
-        quantize_dropout = True
+        quantize_dropout = True,
+        beam_size = 2,
+        eval_beam_size = 3
     )
 
     x = torch.randn(1, 1024, 256)
 
     for _ in range(5):
-        quantized, indices, commit_loss = residual_vq(x, beam_size = 3)
+        quantized, indices, commit_loss = residual_vq(x)
 
     assert quantized.shape == (1, 1024, 256)
     assert indices.shape == (1, 1024, 8)

vector_quantize_pytorch/residual_vq.py

@@ -235,8 +235,10 @@ def __init__(
 
         # beam size
 
-        self.beam_size = default(beam_size, eval_beam_size)
-        self.eval_beam_size = eval_beam_size
+        assert not (exists(eval_beam_size) and not exists(beam_size))
+
+        self.beam_size = beam_size
+        self.eval_beam_size = default(eval_beam_size, beam_size)
 
         # able to assign a different weight for the scoring at each quantizer layer