@avinash_m
A typo in my previous reply, it should be out = torch::log_softmax(out.to(torch::kDouble), 1).to(torch::kLong)
It should be backward compatible. Can I see your back-propogation error log?
The previous error log you showed me is caused by calling log_softmax in forward.