move application of noise when noise dropout is on for fsq

Author: lucidrains

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "vector-quantize-pytorch"
-version = "1.27.16"
+version = "1.27.19"
 description = "Vector Quantization - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

vector_quantize_pytorch/finite_scalar_quantization.py

@@ -159,17 +159,20 @@ def symmetry_preserving_bound(self, z, hard_clamp = False):
     def quantize(self, z):
         """ Quantizes z, returns quantized zhat, same shape as z. """
 
-        shape, device, noise_dropout, preserve_symmetry = z.shape[0], z.device, self.noise_dropout, self.preserve_symmetry
+        shape, device, preserve_symmetry = z.shape[0], z.device, self.preserve_symmetry
         bound_fn = self.symmetry_preserving_bound if preserve_symmetry else self.bound
 
-        bounded_z = bound_fn(z, hard_clamp = self.bound_hard_clamp)
+        return bound_fn(z, hard_clamp = self.bound_hard_clamp)
 
-        # determine where to add a random offset elementwise
-        # if using noise dropout
+    def maybe_apply_noise(self, bounded_z):
+        noise_dropout = self.noise_dropout
 
         if not self.training or noise_dropout == 0.:
             return bounded_z
 
+        # determine where to add a random offset elementwise
+        # if using noise dropout
+
         offset_mask = torch.full_like(bounded_z, noise_dropout).bernoulli_().bool()
         offset = torch.rand_like(bounded_z) - 0.5
 
@@ -270,6 +273,8 @@ def forward(self, z):
             if self.return_indices:
                 indices = self.codes_to_indices(codes)
 
+            codes = self.maybe_apply_noise(codes)
+
             codes = rearrange(codes, 'b n c d -> b n (c d)')
 
             codes = codes.to(orig_dtype)