Loss not converging

Hi,

Thanks for creating this repo! Great Work!

I am trying to run the main.py and main.ipynb file, but during the training, the loss doesn't go down as described in the main.ipynb illustration, and it remains jiggling around 500. Can you let me know what couplings and generator functions you used to get it converged, or used any other thing as well?

Best,
Yogesh