model/mlp.py

import torch.nn as nn


class FeedForward(nn.Module):

    def __init__(

        self,

        d_model,

        dropout=0.1

    ):
        super().__init__()

        self.network = nn.Sequential(

            nn.Linear(
                d_model,
                4 * d_model
            ),

            nn.GELU(),

            nn.Linear(
                4 * d_model,
                d_model
            ),

            nn.Dropout(
                dropout
            )
        )

    def forward(self, x):

        return self.network(x)