Add gradient penalty

[borisdayma]

We need to add a gradient penalty in the GAN discriminator similar to StyleGAN for stability.

Katherine Crowson suggested spectral norm as an alternative and maybe EMA.

[borisdayma]

This is done