VAE pre-trained weights: either I'm doing inference wrong or they are not as good as expected

[RR-28023]

Hi.
I'm trying to reproduce the results shown in the docs with regards to reconstruction quality using the CIFAR-10 pre-trained VAE (LHS are the real images, RHS are the generated)

Here is a link to a self-contained colab notebook that reproduces the steps I follow to perform the inference and reconstruct the pictures. It is a simple script that loads the model, the weights, and tests it over the training set. My problem is that I get a much worse reconstruction (top row are real images, bottom row are the generated ones):

Am I doing something wrong? Could it be that those weights are not as good as one would expect from the image in the docs?
Thanks