* Removing `retain_graph` in backward for discriminator * Fixing issues with different precision training for Lightining