summaryrefslogtreecommitdiff
path: root/text_recognizer/networks/vqvae
diff options
context:
space:
mode:
Diffstat (limited to 'text_recognizer/networks/vqvae')
-rw-r--r--text_recognizer/networks/vqvae/vqvae.py1
1 files changed, 1 insertions, 0 deletions
diff --git a/text_recognizer/networks/vqvae/vqvae.py b/text_recognizer/networks/vqvae/vqvae.py
index e8660c4..f31b062 100644
--- a/text_recognizer/networks/vqvae/vqvae.py
+++ b/text_recognizer/networks/vqvae/vqvae.py
@@ -38,6 +38,7 @@ class VQVAE(nn.Module):
return self.pre_codebook_conv(z_e)
def quantize(self, z_e: Tensor) -> Tuple[Tensor, Tensor]:
+ """Quantizes the encoded latent vectors."""
z_q, vq_loss = self.quantizer(z_e)
return z_q, vq_loss