Skip to content

Commit 20c229e

Browse files
committed
fix: skip round+normalize when preserve_symmetry=True (already in [-1,1])
1 parent 9d2594b commit 20c229e

File tree

1 file changed

+3
-0
lines changed

1 file changed

+3
-0
lines changed

vector_quantize_pytorch/finite_scalar_quantization.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,9 @@ def quantize(self, z):
151151
offset = torch.rand_like(bounded_z) - 0.5
152152
bounded_z = torch.where(offset_mask, bounded_z + offset, bounded_z)
153153

154+
if preserve_symmetry:
155+
return bounded_z
156+
154157
return round_ste(bounded_z) / half_width
155158

156159
def _scale_and_shift(self, zhat_normalized):

0 commit comments

Comments
 (0)