Issue when training VAE

I am trying to train a Variational Autoencoder (VAE) to learn a curve, $f_1(x)=x^2+omega$. It can learn it to a great accuracy, but the problem is that once is trained, for any input to the encoder, the decoder will always output $x^2$. In particular, when training it seems the decoder just learns how to map any input to $x^2$, so the problem seems to be in the decoder, but I am not sure how to solve it. Has anyone experienced this problem? I’d appreciate any suggestion or idea on how to solve this issue.