Getting Nan loss after some time when training custom AnoVAEGAN model.

I am training a custom AnoVAEGAN model using mixed precision and mirror strategy.

I have usedAdam (with epsilon 1e-04), lr = 0.0001,MSE and BinaryCrossEntropy loss for generator.BinaryCrossEntropy for the discriminator.

For some reason, after some time, MSE loss becomes Nan which makes the overall generator loss also Nan.

Any advice from anyone who has also faced a similar issue, please help me out here.

Github repo: https://github.com/dhruvampanchal/AnoVAEGAN

Thanks!

Check comment for nan loss output.