Update convnext

1 files changed, 3 insertions, 6 deletions

diff --git a/text_recognizer/network/convnext/downsample.py b/text_recognizer/network/convnext/downsample.py
index a8a0466..dcc14aa 100644
--- a/text_recognizer/network/convnext/downsample.py
+++ b/text_recognizer/network/convnext/downsample.py
@@ -1,6 +1,4 @@
 """Convnext downsample module."""
-from typing import Tuple
-
 from einops.layers.torch import Rearrange
 from torch import Tensor, nn
 
@@ -8,12 +6,11 @@ from torch import Tensor, nn
 class Downsample(nn.Module):
     """Downsamples feature maps by patches."""
 
-    def __init__(self, dim: int, dim_out: int, factors: Tuple[int, int]) -> None:
+    def __init__(self, dim: int, dim_out: int) -> None:
         super().__init__()
-        s1, s2 = factors
         self.fn = nn.Sequential(
-            Rearrange("b c (h s1) (w s2) -> b (c s1 s2) h w", s1=s1, s2=s2),
-            nn.Conv2d(dim * s1 * s2, dim_out, 1),
+            Rearrange("b c (h s1) (w s2) -> b (c s1 s2) h w", s1=2, s2=2),
+            nn.Conv2d(dim * 4, dim_out, 1),
         )
 
     def forward(self, x: Tensor) -> Tensor: