def __init__(self, out_features: int):
        """
        Args:
            out_features: size of the output tensor
        """
        super().__init__()
        layers = [
            nn.Conv2d(1, 32, 3, 1),
            nn.LeakyReLU(),
            nn.BatchNorm2d(32),
            nn.Conv2d(32, 64, 3, 1),
            nn.LeakyReLU(),
            nn.MaxPool2d(2),
            Flatten(),
            nn.BatchNorm1d(9216),
            nn.Linear(9216, 128),
            nn.LeakyReLU(),
            nn.Linear(128, out_features),
            nn.BatchNorm1d(out_features),
        ]

        self._net = nn.Sequential(*layers)

        Args:
            out_features: size of the output tensor
        
catalyst/contrib/models/mnist.py