Indeed, it would be most elegant. The reason I used a hook was because it was suggested to use those when manually adjusting gradients. I think though for this problem that could be a nice solution.
I’m new to PyTorch so defining custom layers is something I have never done. After running your snippet, could I use the layer by simply treating it as a normal layer? That is:
class domain_classifier(nn.Module):
    def __init__(self):
        super(domain_classifier, self).__init__()
        self.fc1 = nn.Linear(1200, 100) 
        self.fc2 = nn.Linear(100, 1)
    
    def forward(self, x):
        x = grad_reverse(x)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.sigmoid(x)