import torch
from torch import nn


class FlattenToVector(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, inputs):
        return inputs.view(inputs.shape[0], -1)


class TorchMLPClassifier(nn.Module):
    def __init__(self, input_dim, output_dim, hidden_dim=100, negative_slope=0.01):
        super().__init__()
        self.network = nn.Sequential(
            FlattenToVector(),
            nn.Linear(input_dim, hidden_dim),
            nn.LeakyReLU(negative_slope=negative_slope),
            nn.Linear(hidden_dim, output_dim),
        )
        print('pytorch model initialized ...')

    def forward(self, inputs):
        return self.network(inputs)


if __name__ == '__main__':
    model = TorchMLPClassifier(784, 10)
