def __init__(self, input_dim, num_heads, dim_feedforward, dropout=0.0):
                """
                Inputs:
                    input_dim - Dimensionality of the input
                    num_heads - Number of heads to use in the attention block
                    dim_feedforward - Dimensionality of the hidden layer in the MLP
                    dropout - Dropout probability to use in the dropout layers
                """
                super().__init__()
                # Attention layer
                self.attn = torch.nn.MultiheadAttention(
                    embed_dim=input_dim,
                    num_heads=num_heads,
                    batch_first=True,
                )
                # # Two-layer MLP
                self.linear_net = torch.nn.Sequential(
                    torch.nn.Linear(input_dim, dim_feedforward),
                    torch.nn.Dropout(dropout),
                    torch.nn.ReLU(inplace=True),
                    torch.nn.Linear(dim_feedforward, input_dim),
                )
                # Layers to apply in between the main layers
                self.norm1 = torch.nn.LayerNorm(input_dim)
                self.norm2 = torch.nn.LayerNorm(input_dim)
                self.dropout = torch.nn.Dropout(dropout)

                Inputs:
                    input_dim - Dimensionality of the input
                    num_heads - Number of heads to use in the attention block
                    dim_feedforward - Dimensionality of the hidden layer in the MLP
                    dropout - Dropout probability to use in the dropout layers
                
fx2ait/fx2ait/example/01_transformer_model/test_transformer_encoder.py