Flows for a one-dimensional toy problem

[julia-sand]

Thanks for a great package!
I have been working on a toy model but I have been stuck for a while on something which seems like it should be fairly straightforward. I would like to use a normalising flow to learn a double peaked distribution from a single peaked base distribution.
I am training with a forward KL loss using samples

What layers in the library work best in a one-dimensional setting?

I have used AffineConstFlow, which does scale and shift the distribution but I can't get a double peaked distribution. I have also tried planar flows, which diverge to infinity after training for even a short time.

Thank you for any advice