Reverse Gradients in PyTorch

gradient_reversal.py

from torch.autograd import Function


class GradientReversal(Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        output = x
        return output

    @staticmethod
    def backward(ctx, grad_output):
        grad_input = None
        if ctx.needs_input_grad[0]:
            grad_input = grad_output.neg()
        return grad_input

reverse_gradients = GradientReversal.apply