@mailcorahul Thanks; after changing the log_softmax() function with yours, the two cross entropy beam closer but still they are not exactly the same. Is this expected or there is mistake somewhere else?

my custom cross entropy: 2.319404125213623
pytorch cross entroopy: 2.6645867824554443