diff options
Diffstat (limited to 'text_recognizer/networks/vqvae')
| -rw-r--r-- | text_recognizer/networks/vqvae/vqvae.py | 1 |
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 |