summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--text_recognizer/networks/vqvae/decoder.py2
-rw-r--r--text_recognizer/networks/vqvae/encoder.py2
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],