Building a custom loss function in pytorch

I want to make a pixel wise softmax loss function.
That is my predicted value would be a matrix of integers and the target would also be a matrix of integers.
Is there a way in which I can create this loss function in pytorch so that it stores the gradients for backprop.

There are other threads showcasing making custom loss functions. Any of those should get you started.