@albanD Any ideas on how I can keep the effect of the gradient hacking local in a forward pass of a module?
def forward(self, x):
    raw = x @ self.weight.t()
    y = self.s1 * raw
    x.register_hook(lambda grad: grad * self.s2 / self.s1)
    return y

How can I prevent other gradients w.r.t. the tensor passed to x to be affected by this hook? I assume something like
h = x.register_hook(lambda grad: grad * self.s2 / self.s1)
h.remove()

would be something like a NO-OP.