Set train/eval correctly on model

[y0ast]

I think you forgot to set train/eval mode in your train/test function. Since VectorQuantizedVAE contains batchnorm that is required to obtain correct results.

You do seem to set it in the pixelcnn prior training. I didn't check all parts of the code for this setting.