register_backward_hook 18 might nor properly work for “complex” modules as described in the Warning box in the docs.
I would thus recommend to use param.register_hook directly, which will only get the gradient as its input.