Real NVP for Banana shape 2D pdf: val loss diverges

[jecampagne]

Hello,
I'm trying to fit NVP bijector on pdf(x1,x2)=N(x1|(1/4)*x2**2,1)N(x2|0,4) (Papamakarios et al) banana shape pdf.

Notice that I succeed to get NVP working on other distributions (as shown bellow). But here for "banana" density I get the following val_losses

STEP: 0; training loss: 4.672626495361328, validation loss: 5.6602678298950195
STEP: 100; training loss: 3.3686559200286865, validation loss: 4.461033821105957
STEP: 200; training loss: 3.14951229095459, validation loss: 8.736567497253418
STEP: 300; training loss: 3.005864143371582, validation loss: 8.117671966552734
STEP: 400; training loss: 2.991278886795044, validation loss: 7.736321449279785
STEP: 500; training loss: 2.761259078979492, validation loss: 1291.7677001953125
STEP: 600; training loss: 2.898589611053467, validation loss: 6.403262138366699
STEP: 700; training loss: 2.6764931678771973, validation loss: 750117.875
STEP: 800; training loss: 2.640979051589966, validation loss: inf
STEP: 900; training loss: 2.5165188312530518, validation loss: inf
STEP: 1000; training loss: 2.559035062789917, validation loss: 8.277564671749436e+27

Do you have an idea to get it right? Thanks