Remove depth wise conv class

1 files changed, 9 insertions, 4 deletions

diff --git a/text_recognizer/networks/conformer/conv.py b/text_recognizer/networks/conformer/conv.py
index f031dc7..ac13f5d 100644
--- a/text_recognizer/networks/conformer/conv.py
+++ b/text_recognizer/networks/conformer/conv.py
@@ -4,7 +4,6 @@ from einops.layers.torch import Rearrange
 from torch import nn, Tensor
 
 
-from text_recognizer.networks.conformer.depth_wise_conv import DepthwiseConv1D
 from text_recognizer.networks.conformer.glu import GLU
 
 
@@ -21,12 +20,18 @@ class ConformerConv(nn.Module):
         self.layers = nn.Sequential(
             nn.LayerNorm(dim),
             Rearrange("b n c -> b c n"),
-            nn.Conv1D(dim, 2 * inner_dim, 1),
+            nn.Conv1d(dim, 2 * inner_dim, 1),
             GLU(dim=1),
-            DepthwiseConv1D(inner_dim, inner_dim, kernel_size),
+            nn.Conv1d(
+                in_channels=inner_dim,
+                out_channels=inner_dim,
+                kernel_size=kernel_size,
+                groups=inner_dim,
+                padding="same",
+            ),
             nn.BatchNorm1d(inner_dim),
             nn.Mish(inplace=True),
-            nn.Conv1D(inner_dim, dim, 1),
+            nn.Conv1d(inner_dim, dim, 1),
             Rearrange("b c n -> b n c"),
             nn.Dropout(dropout),
         )