avinash_m:

auto loss = torch::nn::functional::cross_entropy(inputs, targets);


@glaringlee I have identified the issue after careful debugging. In this line I should be passing outputs and targets to calculate loss however I was passing inputs. However, the error log is not useful to trace back the error. It should have reported shapes mismatch instead  was reporting about log_softmax.
Thank you for all the patience and time in helping me throughout.