You can use reduce=False (see NLLLoss docs here 43) If you build pytorch from master, to get exactly this behavior. It’s much faster to use this than define your own custom function in python.
Autograd will work on the the function you provided, as long as input is a Variable and the operations done in computing the function use it as such. You should pass in a Variable input and see if the code works.