I am trying to implement GMVAE, by using dataset on MNIST.
Extract class 0 and class 1 from MNIST and train a GMVAE to cluster the training samples into 2 distinct classes (Hence 2 distributions) and regenerate representations from the clusters, different from traditional VAE that fits all data into a single standard normal distribution.
I have a link to the colab notebook but I am not too sure if my implementation is indeed a GMVAE or still a normal VAE. Please advise.
Link : Google Colab