class FeedForward(torch.nn.Module):

    def __init__(self, embed_dim, coef):
        super().__init__()
        self.layers = nn.Sequential(
            torch.nn.Linear(embed_dim, coef*embed_dim),
            torch.nn.ReLU(),
            torch.nn.Linear(coef*n_embed, embed_dim),
            torch.nn.Dropout(dropout)
        )

    def forward(self, x):
        return self.layers(x)