Fix backward hooks Runtime Error

[evcyen]

Description

Fixes backward hooks that return the gradient raising RuntimeError: hook 'hook' has changed the size of value (issue #1160).

Root cause: PyTorch's register_full_backward_hook expects the hook to return either None or a tuple of tensors that replaces grad_input. TransformerLens was passing the user's return value through unchanged. When the user returned a single tensor, PyTorch received a tensor instead of a tuple and raised "hook has changed the size of value."

Fix: In hook_points.py, the wrapper for backward hooks now normalizes the return value before returning to PyTorch. if the user returns a single tensor, we wrap it as (result,). if they already return a one-element tuple, we pass it through.

Tests: Existing test test_backward_hook_runs_successfully still passes. A minimal repro script that returns a bare tensor now runs without error.

Minimal Repro Script:

import torch
from transformer_lens import HookedTransformer

model = HookedTransformer.from_pretrained("gpt2-small")
model.eval()

prompt = "Hello world"
tokens = model.to_tokens(prompt)

def simple_hook(grad, hook):
    print(f"{hook.name}: shape={grad.shape}")
    return grad

model.add_hook("blocks.0.hook_resid_post", simple_hook, dir="bwd")
logits = model(tokens)
loss = logits[0, -1, 0]
loss.backward()

model.reset_hooks()
print("OK: backward completed without error")

Fixes #1160

Type of change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility