def __init__(
        self, shape, nhid=16, n_classes=10, device: Union[str, torch.device] = "cpu"
    ):
        """
        :param shape: Shape of each input sample
        :param nhid: Dimension of latent space of Encoder.
        :param n_classes: Number of classes -
                        defines classification head's dimension
        """
        super(MlpVAE, self).__init__()
        self.dim = nhid
        if device is None:
            device = "cpu"

        self.device = torch.device(device)
        self.encoder = VAEMLPEncoder(shape, latent_dim=128)
        self.calc_mean = MLP([128, nhid], last_activation=False)
        self.calc_logvar = MLP([128, nhid], last_activation=False)
        self.classification = MLP([128, n_classes], last_activation=False)
        self.decoder = VAEMLPDecoder(shape, nhid)

        :param shape: Shape of each input sample
        :param nhid: Dimension of latent space of Encoder.
        :param n_classes: Number of classes -
                        defines classification head's dimension
        
avalanche/models/generator.py