only sync seed during training

lucidrains · Nov 8, 2024 · 35c4c84 · 35c4c84
1 parent cd35a68
commit 35c4c84
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 5 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "vector-quantize-pytorch"
-version = "1.19.2"
+version = "1.19.3"
 description = "Vector Quantization - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }

diff --git a/vector_quantize_pytorch/residual_fsq.py b/vector_quantize_pytorch/residual_fsq.py
@@ -305,7 +305,7 @@ def forward(
 
         forward_kwargs = dict(
             return_all_codes = return_all_codes,
-            rand_quantize_dropout_fixed_seed = get_maybe_sync_seed(device)
+            rand_quantize_dropout_fixed_seed = get_maybe_sync_seed(device) if self.training else None
         )
 
         # invoke residual vq on each group

diff --git a/vector_quantize_pytorch/residual_lfq.py b/vector_quantize_pytorch/residual_lfq.py
@@ -272,7 +272,7 @@ def forward(
         forward_kwargs = dict(
             mask = mask,
             return_all_codes = return_all_codes,
-            rand_quantize_dropout_fixed_seed = get_maybe_sync_seed(device)
+            rand_quantize_dropout_fixed_seed = get_maybe_sync_seed(device) if self.training else None
         )
 
         # invoke residual vq on each group

diff --git a/vector_quantize_pytorch/residual_vq.py b/vector_quantize_pytorch/residual_vq.py
@@ -328,7 +328,6 @@ def forward(
 
             rand = random.Random(rand_quantize_dropout_fixed_seed)
 
-
             rand_quantize_dropout_index = rand.randrange(self.quantize_dropout_cutoff_index, num_quant)
 
             if quant_dropout_multiple_of != 1:
@@ -496,7 +495,7 @@ def forward(
             sample_codebook_temp = sample_codebook_temp,
             mask = mask,
             freeze_codebook = freeze_codebook,
-            rand_quantize_dropout_fixed_seed = get_maybe_sync_seed(device)
+            rand_quantize_dropout_fixed_seed = get_maybe_sync_seed(device) if self.training else None
         )
 
         # invoke residual vq on each group