Thank for reply @albanD,
I have enabled the anomaly detection, otherwise there would be no stacktrace n.1.
After writing this long post after hours of debugging, I have replaced torch Crossentropy loss with my own implementation and now it seems to work  .
loss_s = -F.log_softmax(log_probs_S_per_ex, 0)[gt_start[0]]
loss_e = -F.log_softmax(log_probs_E_per_ex, 0)[gt_end[0]]
loss += loss_s + loss_e

Following your advice, I have also tried changing code to loss = loss + loss_s + loss_e, but results are the same when using torch’s CrossEntropyLoss. Is there an inplace operation somewhere in CrossEntropyLoss loss?