Hi Matej,
check out the Python implementation of F.nll_loss 42 for inspiration.
Of course, if you moved those bits (and even the dispatch to 2d potentially) into nll_loss on the C++ side in PyTorch itself, you could clean up your code. 
Best regards
Thomas