class GradReverse(Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return x

    @staticmethod
    def backward(ctx, grad_output):
        return -grad_output

class Discriminator(nn.Module):
    def __init__(self, channel, pool, classes=10):
        super(Discriminator, self).__init__()
        self.cls_fn = nn.Linear(channel, classes)

    def forward(self, x, reverse=False):
        x = self.cls_fn(x)
        if reverse:
            x = GradReverse.apply(x)
        return x


def train():
    discriminator = Discriminator(3, 10)
    out = discriminator(input, reverse=True)
    ce_loss = CE_criterion(out, label)
    ...
    total_loss += ce_loss

    grads = {}
    def save_grad(name):
       def hook(grad):
          grads[name] = grad
        return hook

    ce_loss.register_hook(save_grad('ce_loss'))
    total_loss.backward()
    print(grads['ce_loss']) # 0.125
   

@albanD
Even if I call Discriminator() with reverse=True, the gradient of ce_loss is still positive value.
What’s wrong with my code ?