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.