VAE not getting trained

Hi,
I am trying to generate samples through VAE, but after training, I can see a blob at the center of the image. I tried mse loss, bce loss but still, the problem persists.
The loss is shooting up and becomes nan values after some epoch.
I tried with mse loss

recons_loss =F.mse_loss(recons, input,reduction='mean')
kld_loss = torch.mean(-0.5 * torch.sum(1 + log_var - mu ** 2 - log_var.exp(), dim = 1), dim = 0)

and also with bce loss

recons_loss = F.binary_cross_entropy_with_logits(recons, input, size_average=False)
        kld_loss = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp()).mean()

Please let me know what is the issue here?

kaggle notebook link: Pokemon | Kaggle