Training the discrete VAE

[Todibo99]

Hi,

I have looked at the usage.ipynb. But I find the training pipeline from the notebook is quite different from what is claimed in the paper. I am wondering how to transform the encoder output to the decoder input. From the paper released I guess the only thing to modify is to replace the argmax in the notebook with a gumbel softmax. Is my understanding correct?

Thanks!