GitLab

* add edm style training

* style

* finish adding edm training feature

* import fix

* fix latents mean

* minor adjustments

* add edm to readme

* style

* fix autocast and scheduler config issues when using edm

* style

---------
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>