I want to have CNN encoder for my texts.
Here is what i have now:
class CNNEncoder(nn.Module):
    def __init__(self, embed_dim: int, vocab_size: int):
        super(CNNEncoder, self).__init__()

        self.embed = nn.Embedding(vocab_size, embed_dim)
        self.conv1 = nn.Conv1d(embed_dim, 120, 1)
        self.pool = nn.AvgPool1d(1, 1)

    def forward(self, x):
        x = self.embed(x)
        if len(x.shape) == 2:
            x = x.unsqueeze_(1).permute([0, 2, 1])
        x = self.conv1(x)
        x = torch.tanh(x)
        x = self.pool(x)
        return x

When i pass it tensor of size [120] it works fine.
But i need to encode a matrix NxK of sentences. It has size [5, 20, 120].
How can i deal with that?