convenience

Author: lucidrains

pyproject.toml

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

vector_quantize_pytorch/finite_scalar_quantization.py

@@ -61,7 +61,7 @@ def floor_ste(z):
 class FSQ(Module):
     def __init__(
         self,
-        levels: list[int],
+        levels: list[int] | tuple[int, ...],
         dim: int | None = None,
         num_codebooks = 1,
         keep_num_codebooks_dim: bool | None = None,
@@ -76,6 +76,9 @@ def __init__(
     ):
         super().__init__()
 
+        if isinstance(levels, tuple):
+            levels = list(levels)
+
         _levels = tensor(levels, dtype = int32)
         self.register_buffer('_levels', _levels, persistent = False)