Hi @albanD
I don’t know about “grad_output” means. Accounding to “dl/dx = dl/dy * dy/dx”, does it means dy / dx?
With my poor understanding:
l = sum(mean( (label - input)**2 ))
dl / dx = -2 * sum(mean( (label - input) )) * dy / dx
Is it correct?
class MSE_SCORE(Function):
    @staticmethod
    def forward(ctx, input, label):
        numpy_input = input.detach().numpy()
        numpy_label = label.detach().numpy()
        diff = (numpy_label - numpy_input)**2
        mean = np.mean(diff, axis=(1,2,3))
        mse = np.sum(mean)
        ctx.save_for_backward(input, label)

        return torch.tensor(mse)

    @staticmethod
    def backward(ctx, grad_output):
        grad_output = grad_output.detach()
        input, label = ctx.saved_tensors
        numpy_input = input.detach().numpy()
        numpy_label = label.detach().numpy()
        grad = -2 * np.sum(np.mean((numpy_label - numpy_input), axis=(1,2,3))) * np.ones(input.shape)
        return grad_output * grad, None