Hi,
As mentionned in the doc for .register_backward_hook(), they are fairly broken right now and will give aberrant behavior :’( We are working on fixing it but it proves quite challenging.
In particular, the grad_input value will be wrong as soon as your nn.Module  does more than one operation.
If you have a single Tensor for input and outputs, you should use tensor.register_hook() that is called on a Tensor in combination with the nn.Module.register_forward_hook() to get the behavior you want:
This would look like this (I did not run the code so there may be typos)
# Add hooks to construct per-example gradients
gradients = {}
activations = {}
def forward_hook(layer, inputs, outputs):
    activations[layer] = inputs[0].detach()

    output_grads = []
    def hook_output(grad_output):
        output_grads[0] = grad_output.detach()

    def hook_input(grad_input):
        A = activations[layer]
        B = output_grads[0]
        gradients[layer] = torch.einsum("ni,nj->nij", B, A)

    inputs[0].register_hook(hook_input)
    outputs[0].register_hook(hook_output)

model.lin.register_forward_hook(forward_hook)