It turns out instead of calling register_backward_hook() on modules, one can check whether they have weight and/or bias properties and call register_hook() on those. The callback will receive the gradient with respect to the tensor.