    def __init__(self, **kwargs):
        """
        Initializes the network
        """
        super().__init__()

        # mnist images are (1, 28, 28) (channels, width, height)
        self.optimizer = None
        self.scheduler = None
        self.layer_1 = torch.nn.Linear(28 * 28, 128)
        self.layer_2 = torch.nn.Linear(128, 256)
        self.layer_3 = torch.nn.Linear(256, 10)
        self.args = kwargs
Initializes the network
tests/integration/frameworks/mlflow/MNIST/mnist_autolog_example.py