I retrained using your code, and after more than ten epochs, the loss turned into nan. Did you encounter this issue during your training as well?