class_docstr: Variational autoencoder module:
  fully-connected and suited for any input shape and type.
  
  The encoder only computes the latent represenations
  and we have then two possible output heads:
  One for the usual output distribution and one for classification.
  The latter is an extension the conventional VAE and incorporates
  a classifier into the network.
  More details can be found in: https://arxiv.org/abs/1809.10635
class_docstr_tok: ["Variational","autoencoder","module",":","fully-connected","and","suited","for","any","input","shape","and","type",".","The","encoder","only","computes","the","latent","represenations","and","we","have","then","two","possible","output","heads",":","One","for","the","usual","output","distribution","and","one","for","classification",".","The","latter","is","an","extension","the","conventional","VAE","and","incorporates","a","classifier","into","the","network",".","More","details","can","be","found","in",":","https",":","\/\/arxiv.org\/abs\/1809.10635"]
docstr: :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
docstr_tok: [":","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"]
code: 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)
  
imports: from abc import abstractmethod
  from typing import Union
  from matplotlib import transforms
  import torch
  import torch.nn
  from torchvision import transforms
  from avalanche.models.utils import MLP, Flatten
  from avalanche.models.base_model import BaseModel
label_desc: Use image node_id 1 to create a new MlpVAE object from inherited base classes: Generator, nn with example: obj = MlpVAE(shape, nhid, n_classes, device)