PyTorch中具有自定义反向函数的损失--简单均方误差示例中的爆炸性损失

import torch

class MyMSELoss(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, target):
        ctx.save_for_backward(input, target)
        loss = torch.mean((input - target)**2)
        return loss

    @staticmethod
    def backward(ctx, grad_output):
        input, target = ctx.saved_tensors
        grad_input = 2 * (input - target) * grad_output
        return grad_input, None

loss_fn = MyMSELoss.apply

output = model(input)
loss = loss_fn(output, target)