Add more residual blocks in vqvae

2 files changed, 2 insertions, 2 deletions

diff --git a/text_recognizer/networks/vqvae/decoder.py b/text_recognizer/networks/vqvae/decoder.py
index 1194573..5279cbd 100644
--- a/text_recognizer/networks/vqvae/decoder.py
+++ b/text_recognizer/networks/vqvae/decoder.py
@@ -31,7 +31,7 @@ class Decoder(nn.Module):
     def _build_decompression_block(self,) -> nn.Sequential:
         in_channels = self.hidden_dim * self.channels_multipliers[0]
         decoder = []
-        for _ in range(2):
+        for _ in range(4):
             decoder += [
                 Residual(
                     in_channels=in_channels,
diff --git a/text_recognizer/networks/vqvae/encoder.py b/text_recognizer/networks/vqvae/encoder.py
index 06e6c18..fe5ef4b 100644
--- a/text_recognizer/networks/vqvae/encoder.py
+++ b/text_recognizer/networks/vqvae/encoder.py
@@ -57,7 +57,7 @@ class Encoder(nn.Module):
                 activation_fn,
             ]
 
-        for _ in range(2):
+        for _ in range(4):
             encoder += [
                 Residual(
                     in_channels=self.hidden_dim * self.channels_multipliers[-1],