Supervised VAE model

Hi team,

Rather than using VAE for reconstruction, I am trying to modify the VAE into a supervised model.

My input dataset consists of images of router locations in an indoor environment and I am predicting the signal strength of the indoor environment from that router location. So my input and label data are separate while training the supervised VAE model. I will later test the VAE model on new indoor environment router locations.

This is what I am trying:
In a VAE model, the reconstruction loss is calculated as,

Recons_loss = BCE (input_data,z_decoded)

I am using Binary cross entropy between the input image and the decoder output image to calculate the reconstruction loss here.

However, I want to calculate the reconstruction loss between the decoder output and signal strength image to better predict the signal strength of the indoor environment.

Kindly let me know if my method is right or if there is any other method to better train a supervised VAE model.

Thanks,
Rahul