Variational autoencoder with probabilistic layers

Hi!
I’m trying to build a variational autoencoder adjusting the tutorial example with probabilistic layers.

My data are grayscale images with pixel values 0 to 255. So I’m confused on how to adjust the decoder code from the example. Particularly Im confused what distribution layer to add to the output of the decoder. I think an probabilistic layer would make for an easy loss function (log_prob).

Also my pixels in most of the pictures follow a similar distribution.

I’m new here and new with deep learning so sorry if I’m asking something trivial.

Thanks in advance.