A simple and clean Multilayer Perceptron implementation (pytorch)

mlp.py

from typing import List
from itertools import chain
import torch
from torch import nn


class MLP(nn.Module):
    def __init__(self,
                hidden_sizes: List[int],
                act_fn: nn.Module = nn.ReLU,
                layer_norm: bool = True,
                bias: bool = True,
                dropout=.2):
        super().__init__()

        self.layers = nn.Sequential(*[
            nn.Sequential(
                nn.Linear(in_features=in_size, out_features=out_size, bias=bias),
                act_fn(),
                nn.LayerNorm(out_size) if layer_norm else nn.Identity(),
                nn.Dropout(dropout)) \
            for in_size, out_size in zip(hidden_sizes[:-1], hidden_sizes[1:])])

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