def METHOD_NAME(ctx, grad_output):
    # Retreive tensors from the forward path.
    softmax, target_mask, masked_target_1d = ctx.saved_tensors
    label_smoothing, vocab_size = ctx.label_smoothing, ctx.vocab_size
    # All the inputs have softmax as thier gradient.
    grad_input = softmax
    # For simplicity, work with the 2D gradient.
    partition_vocab_size = softmax.size()[-1]
    grad_2d = grad_input.view(-1, partition_vocab_size)
    # Add the gradient from matching classes.
    arange_1d = torch.arange(start=0, end=grad_2d.size()[0], device=grad_2d.device)
    softmax_update = 1.0 - target_mask.view(-1).float()
    if label_smoothing > 0:
        smoothing = label_smoothing * vocab_size / (vocab_size - 1)
        grad_2d[arange_1d, masked_target_1d] -= (1.0 - smoothing) * softmax_update
        average_grad = 1 / vocab_size
        grad_2d[arange_1d, :] -= smoothing * average_grad
    else:
        grad_2d[arange_1d, masked_target_1d] -= softmax_update
    # Finally elementwise multiplication with the output gradients.
    grad_input.mul_(grad_output.unsqueeze(dim=-1))
    return grad_input, None, None