When I reproduce your work, the loss value is Nan. There may be a problem with the loss function. I'd like to ask if you have a solution.