Building a custom loss function in pytorch

1

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.

2

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