Hey everyone,
Currently I am developing an multiclass classifikator. I try to replace the nn.CrossEntropyLoss  with my own loss function. The background is that I want to penalize some missclassification stronger than others. For example I have 4 classes, then the following matrix will describe the penalty:
0 , 1 , 4 , 5
1 , 0 , 3 , 8
4 , 3 , 0 , 2
5 , 8 , 2 , 0
As you can see, a misclassification from a sample from class 1 as class 2 is not so fatal as class 4. In the first case, my loss function should return 1 and in the second case 5.  I started to implement the function as follow:
class MyLoss(nn.Module):
    def __init__(self):
         super(MyLoss, self).__init__()
        
    def forward(self, output, target):
           _, predicted = torch.max(output, axis=1)
           loss = []
           for i in range(len(predicted):
                   loss.append(self._loss_matrix[predicted[i].cpu().numpy()][target[i].cpu().numpy()]
           return torch.mean(loss)

I figured out that the loss can not be backpropagated properly since the torch.max function destroys the backpropagation graph. Has anybody any idea how to replace the torch.max function with a differentiable function. I know that I can use softmax(output)[1] for a 2 class classification problem. Is there something similar for the n class problem?
Thanks in advance!!