@ptrblck Sorry for the late reply.
It seems like its the discriminator model thats being affected by NaN. I printed out the values of all the parameters of the discriminator. After 300 epoch, the values in discriminator parameters are in the order of  e-03 to e-05 and have non nan gradients. I trained for a bit and I could not find any gradients being NaN or inf. I think that after the optimizer updates the parameters, the parameter values become so small that it results in NaN.
However when I load the most recent epoch that gave a NaN, I could find the most of the discriminator parameters are NaN. I could also find that some of the gradients were inf while most of the gradients were NaN.
So, is my training loop that i had shared in my original question even correct? In the sense that is the usage of autocast() for discriminator model correct?
The discriminator model gets updated twice in the same iteration of a epoch and then the generator gets updated. So after first discriminator update, I have performed scaler.update(). I have skipped scaler.update() after the second discriminator update. Then finally I have done scaler.update() after generator model gets updated.
Should I add a scaler.update() after the second time discriminator gets updated? Or I could disable autocast() for discriminator model.