self.layers = nn.Sequential(
    # Expand and then reduce the embedding dimensions
    nn.Linear(cfg["emb_dim"], 4 * cfg["emb_dim"]),
    GELU(),
    nn.Linear(4 * cfg["emb_dim"], cfg["emb_dim"]),
)

        __
        __