import torch.nn as nn
from .activation import GELU

class FeedForward(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(cfg['emb_dim'], 4 * cfg['emb_dim']),
            GELU(),
            nn.Linear(4 * cfg['emb_dim'], cfg['emb_dim']),
            nn.Dropout(cfg['drop_rate'])
        )
    
    def forward(self, x):
        return self.layers(x)