VAE protein design

Recently, I’ve embarked on a solo endeavor focused on protein design, leveraging Variational Autoencoders (VAE) for gene therapy advancements. Would anyone be willing to review my work and share their thoughts or suggestions especially in network architecture and the training process? Also anyone willing to coach me on this project?

My main concern is that after 3 epochs my kld goes inf and bce goes nan. How can i get to the root cause of it?

Working version that I hope someone will review.