Cannot reproduce Fig. 4A from the paper

[dpiras]

Hello there,

thank you for the great work! I was trying to reproduce the latent space distribution you show in Fig. 4A in the paper, and in section 2 in this repo. After cloning the repository, downloading the data and setting the paths to data and save folders, I run:

python aae_mnist.py --train --label --dist gmm --ncode 2 --maxepoch 401

and even after 400 epochs what I get in the latent space does not look like your picture, unfortunately. I'm using 10k labelled data, but not much changes if I use the full labels. Any ideas about what I may be missing?

I'm using tensorflow 1.13.1, tensorflow_probability 0.6.0, and training on a NVIDIA Tesla K80 GPU. Any help is really appreciated!

Thank you,
Davide