From 7e23e2ac2abf0f5a4e0d50127a6e873b209a3c99 Mon Sep 17 00:00:00 2001 From: Flavio Schneider Date: Tue, 13 Sep 2022 10:23:30 +0200 Subject: [PATCH] fix: share budget between residual quantizers --- quantizer_pytorch/quantizer.py | 3 +++ setup.py | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/quantizer_pytorch/quantizer.py b/quantizer_pytorch/quantizer.py index 3611ad2..08e2e78 100644 --- a/quantizer_pytorch/quantizer.py +++ b/quantizer_pytorch/quantizer.py @@ -212,11 +212,14 @@ def __init__(self, num_residuals: int, shared_codebook: bool = True, **kwargs): if not shared_codebook: return + # Share both codebooks and total budget first_vq, *rest_vq = self.quantizers codebooks = first_vq.codebooks + budget_ema = first_vq.budget_ema for quantizer in rest_vq: quantizer.codebooks = codebooks + quantizer.budget_ema = budget_ema def from_ids( self, indices: LongTensor, num_residuals: Optional[int] = None diff --git a/setup.py b/setup.py index 233dac6..8bd063d 100644 --- a/setup.py +++ b/setup.py @@ -3,7 +3,7 @@ setup( name="quantizer-pytorch", packages=find_packages(exclude=[]), - version="0.0.16", + version="0.0.17", license="MIT", description="Quantizer - PyTorch", long_description_content_type="text/markdown",