Important for GAN training. If it is possible to do this already, add to #28
Important for GAN training. If it is possible to do this already, add to #28