Hi, I have a problem to implement custom loss using torch.autograd.Function
I know that there is MSELoss() in pytorch.
But I want to implement it via custom loss with image data.
class customMSE(torch.autograd.Function):
    def forward(ctx,output, target):
        ctx.save_for_backward(output, target)

        loss = torch.norm(output-target,p=2) / 2
    
    def backward(ctx, grad_output):
        output, target = ctx.saved_tensors

        return (output - target) * partial_derivative
 

In this case, how can I define partial derivative of output with respect to the network parameters.
Should I have to implement all the chain rule?
Is there any method to get the partial derivative?