from typing import List
import torch


def sequential_nn(in_features: int, num_classes: int, hidden_layers: List[int]):
    """Gets a sequential neural network model

    The sequential model is defined as

      `input => [linear => relu] => output`

    where `[linear => relu]` has sizes defined in
    `hidden_layers`.

    Args:
        in_features: Dimension of the input features
        num_classes: Number of classes
        hidden_layers: Definition of hidden layers

    Returns:
        A sequential model

    """
    hidden_layers_tensors = []
    for i, hidden_layer_size in enumerate(hidden_layers):
        if i == 0:
            in_features_ = in_features
        else:
            in_features_ = hidden_layers[i - 1]

        hidden_layers_tensors.append(
            torch.nn.Linear(in_features_, hidden_layer_size))
        hidden_layers_tensors.append(torch.nn.ReLU())

    output_layer = [
        torch.nn.Linear(hidden_layers[-1], num_classes),
        torch.nn.Softmax(dim=1)
    ]

    layers = hidden_layers_tensors + output_layer

    model = torch.nn.Sequential(*layers)

    return model