import torch


def create_pytorch_model(
    number_of_features, number_of_categories_by_label, label_to_test, **kwargs
):
    """
    Create a toy PyTorch model that has the right shape for inputs and outputs
    """

    class TorchToy(torch.nn.modules.module.Module):
        def __init__(self, number_of_features, number_of_categories):
            super(TorchToy, self).__init__()
            self.fc1 = torch.nn.Linear(number_of_features, 128)
            self.relu1 = torch.nn.ReLU()
            self.fc2 = torch.nn.Linear(128, number_of_categories)

        def forward(self, x):
            x = self.fc1(x)
            x = self.relu1(x)
            x = self.fc2(x)
            return x

    number_of_categories = number_of_categories_by_label[label_to_test]
    model = TorchToy(number_of_features, number_of_categories)
    return model