Fix mixed precision torch compile bug

[Skylion007]

This is solves the bug that torch compile encountered. It didn't know the module was in fp16 until it hit weights in the backward pass because it does not know that .half() has side-effects. Making the program adhere to a more functional paradigm completely solves this bug (by invalidating the dtype cache of the nn Module).

[mvpatel2000]

Nice!!

[jon-chuang]

@Skylion007 could you give more context on:

1. Where the torch.compile was invoked?

[jon-chuang]

Something that I find extremely suspicious about this fix is the following:

self.text_encoder.requires_grad_(False)
self.vae.requires_grad_(False)
if self.encode_latents_in_fp16:
    self.text_encoder = self.text_encoder.half()
    self.vae = self.vae.half()

Since vae does not require grad, how can it throw an error:

RuntimeError: attempting to assign a gradient with dtype 'float' to a tensor with dtype 'c10::Half'. Please ensure that the gradient and the tensor have the same dtype

About the dtype of its grad?