Thanks for sharing the code. I test the nll_loss function and get:
RuntimeError: Input tensor must have same dimensions as output tensor at /b/wheel/pytorch-src/torch/lib/TH/generic/THTensorMath.c:424
It seems that the input and index arguments of gather function should have same dimensions.