address #209

Author: lucidrains

vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -974,6 +974,9 @@ def __init__(
         # for variable lengthed sequences, whether to take care of masking out the padding to 0 (or return the original input)
         self.return_zeros_for_masked_padding = return_zeros_for_masked_padding
 
+        # whether to freeze the codebook, can be overridden on forward
+        self.freeze_codebook = freeze_codebook
+
     @property
     def codebook(self):
         codebook = self._codebook.embed
@@ -1042,12 +1045,16 @@ def forward(
         mask = None,
         lens = None,
         sample_codebook_temp = None,
-        freeze_codebook = False,
+        freeze_codebook = None,
         return_loss_breakdown = False,
         codebook_transform_fn: Callable | None = None
     ):
         orig_input, input_requires_grad = x, x.requires_grad
 
+        # freezing codebook
+
+        freeze_codebook = default(freeze_codebook, self.freeze_codebook)
+
         # handle masking, either passed in as `mask` or `lens`
 
         assert not (exists(mask) and exists(lens))