Lite transformer encoder layer

lite_transformer.py

class LiteTransformerLayer(nn.Module):
    def __init__(self, dim=496, heads=4, kernel_size=4):
        super().__init__()
        assert dim % 2 == 0

        self.attention = Attention(dim // 2, heads=heads)
        self.cnn = LightweightConv(dim // 2, kernel=kernel_size)  # or Dynamic conv

        self.fc = nn.Sequential(
            nn.Linear(dim, dim),
            nn.ReLU(),  # chose your favorite nonlinearity here
            nn.Linear(dim, dim),
        )
        self.norm = nn.LayerNorm(dim)
    
    def forward(self, x):
        _, _, dim = x.shape
        residual = x

        x1 = self.attention(x[:dim // 2])
        x2 = self.cnn(x[dim // 2:])

        x = torch.cat([x1, x2], dim=-1)
        x = self.norm(x + residual)

        residual = x
        x = self.fc(x)
        x = self.norm(x + residual)
        return x